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

OK, just to clarify here, lets do a simple example

A=diag(arange(0,10,1))

gives

array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 4, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 7, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 8, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 9]])

then

A[0][0:4]

gives

array([0, 0, 0, 0])

that is first row, elements 0 to 3. But

A[0:4][1]

doesn't give the first 4 rows, the 2nd element in each. Instead we get

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

i.e the entire 2nd column.

A[0:4,1]

gives

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

I'm sure there is a very good reason for this and which makes perfect sense to programmers but for those of us uninitiated in that great religion it can be quite confusing.