如何读取索引列表并返回多维numpy数组的值

Question

给定一个N维正方形(立方体等(numpy数组n和一个索引列表l，我希望函数返回n[l[0]][l[1]]...[l[N-1]]

例如：

如果

n = [[0,1,2],
[3,4,5],
[6,7,8]]

和

l = [1,1]

操作

f(n,l) = 4

或者如果

n = [[[0,1],
[2,3]],
[[4,5],
[6,7]]]

和

l = [1,0,1]

操作

f(n,l) = 5

如有任何帮助，我们将不胜感激。我正在尝试制作一个代码，在这里我可以改变模拟的维度，而不必为不同的维度重新编码一些新的东西。

Answer 1

在我看来，您只需要numpy就可以完成这项工作。使用n的示例，您可以进行

n =np.array( [[0,1,2],
[3,4,5],
[6,7,8]])
l = (1,1)
n[l]

注意，l需要是一个元组(圆括号(。您可以将您的列表转换为

l = [1,1]
n[tuple(l)]

Answer 2

只需将数组转换为元组：

import numpy as np

def f(n, i):
return n[tuple(i)]

print(
f(
np.arange(0, 9).reshape((3, 3)),
np.array([1, 1])
)
)
print(
f(
np.arange(0, 8).reshape((2, 2, 2)),
np.array([1, 0, 1])
)
)

输出：

4
5

请参阅：将元组/数组/列表解压缩为Numpy arrays 的索引