Numpy - 用索引数组索引数组的最后一个维度答案

【问题标题】：Numpy - Index last dimension of array with index arrayNumpy - 用索引数组索引数组的最后一个维度
【发布时间】：2019-11-20 23:08:40
【问题描述】：

我正在尝试使用由我希望保留的索引组成的矩阵来索引 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) 矩阵，所以我不完全确定我应该如何处理这个问题。

【问题讨论】：

标签： python arrays numpy

【解决方案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]

【讨论】：

numpy has no attribute 'take_along_axis'
@LucEvertzen np.take_along_axis 在最近的 NumPy 版本中被引入。更新 NumPy 或使用 this solution 作为替代或使用第二种方法。

【解决方案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 的形状相匹配。

【讨论】：