按列索引切片numpy.ndarray

时间:2014-10-03 08:46:14

标签: python numpy

有没有办法在不必定义行索引的情况下对下面的数组进行切片,即不必编写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]切片,但看起来不是。

3 个答案:

答案 0 :(得分:3)

您可能正在寻找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,))

答案 1 :(得分:1)

您使用diag(或diagonal)获取X[:,L]的对角元素:

np.diag(X[:,L])

另一种方法是使用where

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

答案 2 :(得分:0)

请注意

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