Question

是否有一种简单的方法可以使用索引数组在最后一个维度上索引numpy多维数组？例如，采用形状a的数组(10, 10, 20)。假设我有一个形状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]]不同。等等。有没有一种简单的方法可以在不诉诸循环的情况下做到这一点？

Answer 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

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

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

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

1 个答案: