沿最后一维索引numpy nd数组 [英] Index numpy nd array along last dimension

问题描述

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

Is there an easy way to index a numpy multidimensional array along the last dimension, using an array of indices? For example, take an array a of shape (10, 10, 20). Let's assume I have an array of indices b, of shape (10, 10) so that the result would be 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]]不同.等等.有没有一种简单的方法可以执行此操作而不求助于循环?

However, this doesn't work because it then tries to index like a[b[i, j], b[i, j]], which is not the same as a[i, j, b[i, j]]. And so on. Is there an easy way to do this without resorting to a loop?

推荐答案

有几种方法可以做到这一点.让我们首先生成一些测试数据:

There are several ways to do this. Let's first generate some test data:

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

解决问题的一种方法是使用

One way to solve the question would be to create two index vectors, where one is a row vector and the other a column vector of 0..9 using meshgrid:

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

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

之所以起作用，是因为i0，i1和b都将广播到10x10矩阵.快速测试正确性:

This works because i0, i1 and b will all be broadcasted to 10x10 matrices. Quick test for correctness:

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 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数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！