[英]Index numpy nd array along last dimension
有没有一种简单的方法可以使用索引数组沿最后一个维索引一个numpy多维数组? 例如,采用形状为(10, 10, 20)
10,10,20)的数组a
。 假设我有一个形状为(10, 10)
的索引b
数组,因此结果为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[b[i, j], b[i, j]]
索引,这与a[i, j, b[i, j]]
。 等等。 有没有一种简单的方法可以执行此操作而不求助于循环?
有几种方法可以做到这一点。 让我们首先生成一些测试数据:
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
,因为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 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
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.