Selecting specific column in each row from array

Question

I am trying to select specific column elements for each row of a numpy array. For example, in the following example:

In [1]: a = np.random.random((3,2))
Out[1]: 
array([[ 0.75670668,  0.1283942 ],
       [ 0.51326555,  0.59378083],
       [ 0.03219789,  0.53612603]])

I would like to select the first element of the first row, the second element of the second row, and the first element of the third row. So I tried to do the following:

In [2]: b = np.array([0,1,0])

In [3]: a[:,b]

But this produces the following output:

Out[3]: 
array([[ 0.75670668,  0.1283942 ,  0.75670668],
       [ 0.51326555,  0.59378083,  0.51326555],
       [ 0.03219789,  0.53612603,  0.03219789]])

which clearly is not what I am looking for. Is there an easy way to do what I would like to do without using loops?

Answer 1

You can use:

a[np.arange(3), (0,1,0)]

in your example above.

Answer 2

This isn't an answer so much as an attempt to document this a bit. For the answer above, we would have:

>>> import numpy as np
>>> A = np.array(range(6))
>>> A
array([0, 1, 2, 3, 4, 5])
>>> A.shape = (3,2)
>>> A
array([[0, 1],
       [2, 3],
       [4, 5]])
>>> A[(0,1,2),(0,1,0)]
array([0, 3, 4])

Specifying a list (or tuple) of individual row and column coordinates allows fancy indexing of the array. The first example in the comment looks similar at first, but the indices are slices. They don't extend over the whole range, and the shape of the array that is returned is different:

>>> A[0:2,0:2]
array([[0, 1],
       [2, 3]])

For the second example in the comment

>>> A[[0,1],[0,1]]
array([0, 3])

So it seems that slices are different, but except for that, regardless of how indices are constructed, you can specify a tuple or list of (x-values, y-values), and recover those specific elements from the array.