Numpy - 使用索引数组索引数组的最后一维

Question

我试图用一个由我希望保留的索引组成的矩阵来索引一个3D矩阵的最后一个维度。你知道吗

我有一个带有形状的推力值矩阵：

(3, 3, 5)

我想过滤最后一个索引，根据一些标准，使它从大小5减少到大小1。我已经在最后一个维度中找到了符合我标准的索引：

[[0 0 1]
 [0 0 1]
 [1 4 4]]

我想要实现的是：对于第一行和第一列，我想要最后一个维度的第0个索引。对于第一行和第三列，我需要最后一个维度的第一个索引。就保持最终矩阵的索引而言，它将变成(3, 3)2D矩阵，如下所示：

[[0,0,0], [0,1,0], [0,2,1];
 [1,0,0], [1,1,0], [1,2,1];
 [2,0,1], [2,1,4], [2,2,4]]

我很有信心numpy能做到这一点，但我不知道具体怎么做。我不想构建一个嵌套for循环的构造。你知道吗

我已经试过了：

minValidPower = totalPower[:, :, tuple(indexMatrix)]

但是这会产生一个(3, 3, 3, 3)矩阵，所以我不能完全确定应该如何处理这个问题。你知道吗

Answer 1

以a作为输入数组，idx作为索引数组-

np.take_along_axis(a,idx[...,None],axis=-1)[...,0]

或者，对于开放网格-

I,J = np.ogrid[:idx.shape[0],:idx.shape[1]]
out = a[I,J,idx]