slice operation for numpy in python2

Question

If I have a numpy array X:

array([[ 0.13263767,  0.23149757,  0.57097612],
       [ 0.49629958,  0.67507182,  0.6758823 ]])

And an index array Y:

array([1, 2, 1])

I could use X[0:,Y] to index the first row of X, and it will output the:

array([ 0.23149757,  0.57097612,  0.23149757])

my question is, if I have an index array Z with 2 dimensions:

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

I would like to use first row to of Z to index the first row of X, and second row of Z to index the second row of X (Z and X have the same rows). So one way is to use command as follow:

Row_0 = X[0:, Z[0]]
Row_1 = X[1:, Z[1]]

I was wondering if there is a simple way to do this. Thanks

Answer 1

You can use fancy indexing to achieve that:

>>> X[[[0], [1]], Z]
array([[ 0.23149757,  0.57097612,  0.23149757],
       [ 0.49629958,  0.67507182,  0.6758823 ]])

The trick is that the array indexing the first dimension must broadcast with the one indexing the second one. In this case:

>>> np.array([[0], [1]]).shape
(2, 1)
>>> Z.shape
(2, 3)

So the return will be of the broadcast shape, (2, 3) with indices taken from the first array for the first dimension, and from the second array for the second dimension.

For more general cases, you can get the same result as:

>>> X[np.arange(Z.shape[0])[:, None], Z]
array([[ 0.23149757,  0.57097612,  0.23149757],
       [ 0.49629958,  0.67507182,  0.6758823 ]])