沿最后一个维度索引 numpy nd 数组答案

【问题标题】：Index numpy nd array along last dimension沿最后一个维度索引 numpy nd 数组
【发布时间】：2015-02-01 08:12:12
【问题描述】：

是否有一种简单的方法可以使用索引数组沿最后一维索引 numpy 多维数组？例如，取一个形状为(10, 10, 20) 的数组a。假设我有一个索引数组b，形状为(10, 10)，因此结果将是c[i, j] = a[i, j, b[i, j]]。

我尝试了以下示例：

a = np.ones((10, 10, 20))
b = np.tile(np.arange(10) + 10, (10, 1))
c = a[b]

但是，这不起作用，因为它会尝试像a[b[i, j], b[i, j]] 那样进行索引，这与a[i, j, b[i, j]] 不同。等等。有没有不使用循环的简单方法来做到这一点？

【问题讨论】：

只是为了确保我理解正确，你想要c[i, j] = a[i, j, b[i, j]]
@mgilson 是的。应该更清楚地说明这一点。
这通常像c = a[np.arange(b.shape[0]), np.arange(b.shape[1]), b] 一样完成，但我希望有更好的方法。
@moarningsun 我的正确性测试失败了。我认为您必须将两个范围之一转换为列向量左右。

标签： python numpy

【解决方案1】：

有几种方法可以做到这一点。我们先生成一些测试数据：

In [1]: a = np.random.rand(10, 10, 20)

In [2]: b = np.random.randint(20, size=(10,10))  # random integers in range 0..19

解决问题的一种方法是创建两个索引向量，其中一个是行向量，另一个是 0..9 的列向量，使用 meshgrid：

In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True)

In [4]: c = a[i0, i1, b]

这是有效的，因为i0、i1 和 b 都将被广播到 10x10 矩阵。快速测试正确性：

In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[5]: True

另一种方法是使用choose 和rollaxis：

# choose needs a sequence of length 20, so move last axis to front
In [22]: aa = np.rollaxis(a, -1)  

In [23]: c = np.choose(b, aa)

In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[24]: True

【讨论】：

谢谢。我更喜欢你的第一个选择。 choose 选项看起来不错，但它不够通用，因为 choose 不适用于大于 32 的索引（numpy 中的当前错误，请参阅 discussion on increasing NPY_MAXARGS）。根据我的测试，它的速度也是meshgrid 选项的两倍。
在这种情况下，您可能应该查看mgrid 和ogrid 之类的函数，这些函数引用自meshgrid 文档。但是您是否使用超过 32 维的矩阵？我为你感到难过，当我试图理解一个超过 2 维的数组时，我的脑袋已经爆炸了:)。
啊，我现在从 10x10x20 变为 10x10x40 时看到了您的错误。我猜choose 将 3D 矩阵解包为沿第一维的 2D 矩阵列表。
@BasSwinckels 我尝试了不相等的行和列大小，但它不起作用（IndexError: shape mismatch: indexing arrays could not be broadcast together）。我对meshgrid不熟悉，你能更新一下吗？