[英]Index 2D numpy array by a 2D array of indices without loops
我正在寻找一种矢量化的方法来通过numpy.array
索引索引numpy.array
。
例如:
import numpy as np
a = np.array([[0,3,4],
[5,6,0],
[0,1,9]])
inds = np.array([[0,1],
[1,2],
[0,2]])
我想建立一个新的阵列,使得阵列中的每一行(i)是阵列中的一行(一) a
,通过阵列INDS(I)的行索引。 我想要的输出是:
array([[ 0., 3.], # a[0][:,[0,1]]
[ 6., 0.], # a[1][:,[1,2]]
[ 0., 9.]]) # a[2][:,[0,2]]
我可以用一个循环来实现这一点:
def loop_way(my_array, my_indices):
new_array = np.empty(my_indices.shape)
for i in xrange(len(my_indices)):
new_array[i, :] = my_array[i][:, my_indices[i]]
return new_array
但我正在寻找一个纯矢量化的解决方案。
当使用索引数组来索引另一个数组时,每个索引数组的形状应该与输出数组的形状相匹配。 您希望列索引匹配inds
,并且您希望行索引匹配输出的行,例如:
array([[0, 0],
[1, 1],
[2, 2]])
你可以只使用上述的单个列中,由于广播,所以可以使用np.arange(3)[:,None]
为垂直arange
因为None
插入一个新的轴:
>>> np.arange(3)[:, None]
array([[0],
[1],
[2]])
最后,一起:
>>> a[np.arange(3)[:,None], inds]
array([[0, 3], # a[0,[0,1]]
[6, 0], # a[1,[1,2]]
[0, 9]]) # a[2,[0,2]]
这样做是有可能的,虽然有点不明显,如下所示:
>>> a[np.arange(a.shape[0])[:, None], inds]
array([[0, 3],
[6, 0],
[0, 9]])
索引np.arange(a.shape[0])
只是索引列索引inds
数组所适用的行。 附加[:, None]
修改此数组的形状,使其形状为(a.shape[0], 1)
,即每个行索引位于 1 列宽的 2D 数组的单独行中。
基本原则是索引数组的维数必须一致,它们的形状也必须一致。 请参阅np.ix_
文档以了解这一点。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.