向量化numpy中的这个for循环

Question

我想知道如何将其向量化为for循环。给定一个2x2x2数组x和一个数组，其中每个元素都是该数组的ith, jth, and kth元素，我想得到x[i,j,k]

给出数组x和y

x = np.arange(8).reshape((2, 2, 2))
y = [[0, 1, 1], [1, 1, 0]]

我想得到：

x[0, 1, 1] = 3和x[1, 1, 0] = 6

我尝试过：

print(x[y]) 

但它会打印：

array([[2, 3],
       [6, 7],
       [4, 5]])

所以我最终做了：

for y_ in y:
    print(x[y_[0], y_[1], y_[2]])

哪个可行，但我不禁想到还有更好的方法。

Answer 1

使用转置的y即zip(*y)作为索引；您需要将每个维度的索引作为advanced indexing的元素才能工作：

x[tuple(zip(*y))]
# array([3, 6])