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.