NumPy使用索引列表选择每行的特定列索引

Question

我很难选择NumPy矩阵每行的特定列.

假设我有以下矩阵,我会称之为X:

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

我还会list调用每行的列索引Y:

[1, 0, 2]

我需要获取值:

[2]
[4]
[9]

取而代之的是list使用索引Y,我还可以产生一个矩阵具有相同的形状X,每一个列是bool/ int范围在0-1之间的值,表明这是否是必需的列.

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

我知道这可以通过迭代数组并选择我需要的列值来完成.但是,这将在大数据阵列上频繁执行,这就是它必须尽可能快地运行的原因.

我想知道是否有更好的解决方案？

谢谢.

Answer 1

如果你有一个布尔数组,你可以根据它做直接选择:

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

与您的初始示例一起,您可以执行以下操作:

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

你也可以添加arange并直接选择,虽然取决于你如何生成布尔数组以及你的代码看起来像YMMV.

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

希望有所帮助,如果您有任何疑问,请告诉我.

Answer 2

你可以这样做:

In [7]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [8]: lst = [1, 0, 2]

In [9]: a[np.arange(len(a)), lst]
Out[9]: array([2, 4, 9])

有关索引多维数组的更多信息:http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Answer 3

最近的numpy版本添加了一个take_along_axis(和put_along_axis)，可以干净地进行索引。

In [101]: a = np.arange(1,10).reshape(3,3)                                                             
In [102]: b = np.array([1,0,2])                                                                        
In [103]: np.take_along_axis(a, b[:,None], axis=1)                                                     
Out[103]: 
array([[2],
       [4],
       [9]])

它的运作方式与：

In [104]: a[np.arange(3), b]                                                                           
Out[104]: array([2, 4, 9])

但具有不同的轴处理。这是特别针对应用的结果argsort和argmax。

Answer 4

一种简单的方法可能如下所示：

In [1]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [2]: y = [1, 0, 2]  #list of indices we want to select from matrix 'a'

range(a.shape[0]) 将返回 array([0, 1, 2])

In [3]: a[range(a.shape[0]), y] #we're selecting y indices from every row
Out[3]: array([2, 4, 9])