Selecting a specific column in each row from an 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 leads to the following conclusion:

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

which is clearly not what I'm 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 the above example.

Answer 2

OK, just for clarification here, let's make 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]) 

this is the first line, elements from 0 to 3. But

 A[0:4][1] 

does not give the first 4 lines, the 2nd element in each. Instead, we get

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

ie the entire second column.

 A[0:4,1] 

gives

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

I am sure that there is a very good reason for this and which makes sense for programmers, but for those of us who are uninitiated in this great religion, this can be quite confusing.

Answer 3

This is not an answer, but an attempt to document a little. 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 column and column coordinates allows for fantastic indexing of the array. The first example in a comment looks similar at first, but indexes are slices. They do not apply to the entire range, and the shape of the returned array 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 the slices are different, but with the exception of this, no matter how the indices are constructed, you can specify a tuple or list (x, y value) and restore these specific elements from the array.