由多维数组中的索引矩阵指定的 numpy 选择矩阵

Question

我有一个大小为 5x5x4x5x5 的 numpy 数组a。我还有另一个大小为 5x5 的矩阵b.我想从 0 到 4 获得i a[i,j,b[i,j]]，从 0 到 4 获得j。这将给我一个5x5x1x5x5矩阵。有没有办法在不使用 2 个for循环的情况下做到这一点？

Answer 1

让我们将矩阵a视为 100 个大小为 (5, 5) 的(= 5 x 5 x 4)矩阵。所以，如果你能得到每个三胞胎的衬里索引 - (i, j, b[i, j]) - 你就完成了。这就是np.ravel_multi_index的用武之地。以下是代码。

import numpy as np
import itertools
# create some matrices
a = np.random.randint(0, 10, (5, 5, 4, 5, 5))
b = np.random(0, 4, (5, 5))
# creating all possible triplets - (ind1, ind2, ind3)
inds = list(itertools.product(range(5), range(5)))
(ind1, ind2), ind3 = zip(*inds), b.flatten()
allInds = np.array([ind1, ind2, ind3])
linearInds = np.ravel_multi_index(allInds, (5,5,4))
# reshaping the input array
a_reshaped = np.reshape(a, (100, 5, 5))
# selecting the appropriate indices
res1 = a_reshaped[linearInds, :, :]
# reshaping back into desired shape
res1 = np.reshape(res1, (5, 5, 1, 5, 5))
# verifying with the brute force method
res2 = np.empty((5, 5, 1, 5, 5))
for i in range(5):
    for j in range(5):
        res2[i, j, 0] = a[i, j, b[i, j], :, :]
print np.all(res1 == res2)  # should print True

Answer 2

有

np.take_along_axis正是为此目的 -

np.take_along_axis(a,b[:,:,None,None,None],axis=2)