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

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

假设我有下面这个矩阵，我称之为 X:

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

我还有一个每行列索引的 list，我称之为 Y:

[1, 0, 2]

我需要得到数值:

[2]
[4]
[9]

与使用 Y索引的 list不同，我还可以生成一个与 X形状相同的矩阵，其中每一列都是0-1范围内的 bool/int，表明这是否是所需的列。

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

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

因此，我想知道是否有更好的解决办法？

67477 次浏览
小开
最佳答案

如果你有一个布尔数组，你可以像这样直接选择:

>>> 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])
小开

你可以这样做:

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

小开

你可以用迭代器来做，像这样:

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
小开

一个简单的方法可能看起来像:

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])
小开

另一个聪明的方法是首先转置数组，然后对其进行索引。最后，选择对角线，它总是正确的答案。

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]
小开

最近的 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测试结果的应用。

小开

使用 take _ along _ axis 从 hpaulj 回答问题应该是被接受的。

下面是一个带有 N-dim 索引数组的派生版本:

>>> 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，但插值公式是否仍然看起来不错？


提交答案