Question

我正在尝试使用包含我希望保留的索引的矩阵来索引3D矩阵的最后一个维度。

我有一个具有以下形状的推力值矩阵：

(3, 3, 5)

我想根据某些条件过滤最后一个索引，以便将其从大小5减小到大小1。我已经找到了最后一个维度中符合我条件的索引：

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

我要实现的目标：对于第一行和第一列，我想要最后一个维度的第0个索引。对于第一行和第三列，我想要最后一个维度的第一个索引。就索引而言，最终矩阵将变为(3, 3)二维矩阵，如下所示：

[[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]

Answer 2

您可以为前两个维构建相应的索引数组。这些基本上是：

[0 1 2]
[0 1 2]
[0 1 2]

[0 0 0] 
[1 1 1] 
[2 2 2]

您可以使用meshgrid函数构造它们。我在示例中将它们存储为m1和m2：

vals = np.arange(3*3*5).reshape(3, 3, 5) # test sample
m1, m2 = np.meshgrid(range(3), range(3), indexing='ij')
m3 = np.array([[0, 0, 1], 0, 0, 1], [1, 4, 4]])
sel_vals = vals[m1, m2, m3]

结果的形状与索引数组m1，m2和m3的形状匹配。

numpy-使用索引数组索引数组的最后一个维度

2 个答案: