按列索引切片numpy.ndarray

Question

Is there a way to slice the array below without having to define the row indices ie not having to write range(len(X)) ?

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

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

I thought it was possible to slice with X[:,L] but looks like it's not.

Answer 1

You're probably looking for 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])

Performance note : while this solution is a direct answer to the question, it's quickly becomes not the most optimal with increase of len(X) . As of numpy 1.9.0, np.arange approach is faster:

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

You take the diagonal elements of X[:,L] using diag (or diagonal ):

np.diag(X[:,L])

Another way to do it is with where :

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

Answer 3

Note that

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

And you want the diagonal elements:

So just do:

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