按列索引切片numpy.ndarray

Question

有沒有一種方法可以切片下面的數組而不必定義行索引，即不必寫range(len(X)) ？

X = np.arange(10*2).reshape((10,2))
L = np.random.randint(0,2,10)

Xs = X[range(len(X)),L]

我以為可以用X[:,L]切片，但看起來好像不是。

Answer 1

您可能正在尋找np.choose ：

In [25]: X = np.arange(10*2).reshape((10,2)); X
Out[25]: 
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11],
       [12, 13],
       [14, 15],
       [16, 17],
       [18, 19]])

In [26]: L = np.random.randint(0,2,10); L
Out[26]: array([1, 1, 1, 1, 1, 0, 0, 0, 0, 1])

In [27]: L.choose(X.T)
Out[27]: array([ 1,  3,  5,  7,  9, 10, 12, 14, 16, 19])

In [28]: # or otherwise

In [29]: np.choose(L, X.T)
Out[29]: array([ 1,  3,  5,  7,  9, 10, 12, 14, 16, 19])

性能說明 ：盡管此解決方案是對問題的直接答案，但隨着len(X)增加，它很快就變得不是最佳的。 從numpy 1.9.0開始， np.arange方法更快：

In [17]: %timeit X[range(len(X)), L]
1000 loops, best of 3: 629 µs per loop

In [18]: %timeit X[np.arange(len(X)), L]
10000 loops, best of 3: 78.8 µs per loop

In [19]: %timeit L.choose(X.T)
10000 loops, best of 3: 146 µs per loop

In [20]: X.shape, L.shape
Out[20]: ((10000, 2), (10000,))

Answer 2

您可以使用diag （或diagonal ）獲取X[:,L]的對角線元素：

np.diag(X[:,L])

另一種方法是在where ：

np.where(L,X[:,1],X[:,0])

Answer 3

注意

In [9]: X[:, L]
Out[9]:
array([[ 1,  1,  0,  0,  1,  0,  1,  0,  1,  0],
   [ 3,  3,  2,  2,  3,  2,  3,  2,  3,  2],
   [ 5,  5,  4,  4,  5,  4,  5,  4,  5,  4],
   [ 7,  7,  6,  6,  7,  6,  7,  6,  7,  6],
   [ 9,  9,  8,  8,  9,  8,  9,  8,  9,  8],
   [11, 11, 10, 10, 11, 10, 11, 10, 11, 10],
   [13, 13, 12, 12, 13, 12, 13, 12, 13, 12],
   [15, 15, 14, 14, 15, 14, 15, 14, 15, 14],
   [17, 17, 16, 16, 17, 16, 17, 16, 17, 16],
   [19, 19, 18, 18, 19, 18, 19, 18, 19, 18]])

您想要對角線元素：

所以做：

In [14]: X[:, L].diagonal()
Out[14]: array([ 1,  3,  4,  6,  9, 10, 13, 14, 17, 18])