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

问题描述：

我正在努力选择 NumPy 矩阵每行的特定列。

假设我有以下矩阵，我将其称为X：

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

我list每行都有一个列索引，我称之为Y：

[1, 0, 2]

我需要获取以下值：

[2]
[4]
[9]

除了带有list索引的以外Y，我还可以生成一个具有相同形状的矩阵，X其中每一列都是一个介于 0-1 值之间的bool/ ，表示这是否是所需的列。int

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

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

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

解决方案 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])

解决方案 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-Dimension-arrays

解决方案 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。

解决方案 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])

解决方案 5：

你可以使用迭代器来实现。像这样：

np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)

时间：

N = 1000
X = np.zeros(shape=(N, N))
Y = np.arange(N)

#@Aशwini चhaudhary
%timeit X[np.arange(len(X)), Y]
10000 loops, best of 3: 30.7 us per loop

#mine
%timeit np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)
1000 loops, best of 3: 1.15 ms per loop

#mine
%timeit np.diag(X.T[Y])
10 loops, best of 3: 20.8 ms per loop

解决方案 6：

hpaulj 使用 take_along_axis 给出的答案应该是可以接受的。

这是一个具有 N 维索引数组的派生版本：

>>> arr = np.arange(20).reshape((2,2,5))
>>> idx = np.array([[1,0],[2,4]])
>>> np.take_along_axis(arr, idx[...,None], axis=-1)
array([[[ 1],
        [ 5]],

       [[12],
        [19]]])

请注意，选择操作与形状无关。我用它来通过拟合抛物线来改进可能的矢量值argmax结果：histogram

def interpol(arr):
    i = np.argmax(arr, axis=-1)
    a = lambda Δ: np.squeeze(np.take_along_axis(arr, i[...,None]+Δ, axis=-1), axis=-1)
    frac = .5*(a(1) - a(-1)) / (2*a(0) - a(-1) - a(1)) # |frac| < 0.5
    return i + frac

注意，squeeze删除大小为 1 的维度，从而得到相同形状的i和frac，即峰值位置的整数和小数部分。

我很确定可以避免lambda，但插值公式看起来仍然不错吗？

解决方案 7：

另一个巧妙的方法是先转置数组，然后对其进行索引。最后取对角线，这总是正确的答案。

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
Y = np.array([1, 0, 2, 2])

np.diag(X.T[Y])

步步：

原始数组：

>>> X
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

>>> Y
array([1, 0, 2, 2])

转置以便能够正确索引。

>>> X.T
array([[ 1,  4,  7, 10],
       [ 2,  5,  8, 11],
       [ 3,  6,  9, 12]])

按 Y 顺序获取行。

>>> X.T[Y]
array([[ 2,  5,  8, 11],
       [ 1,  4,  7, 10],
       [ 3,  6,  9, 12],
       [ 3,  6,  9, 12]])

对角线现在应该变得清晰了。

>>> np.diag(X.T[Y])
array([ 2,  4,  9, 12]