CODEWITHSUNDEEP
pythonnumpy
astrofrog
astrofrog

Reputation: 34091

Selecting specific column in each row from array

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?

Upvotes: 6

Views: 11837

Answers (3)

Carlos Santos
Carlos Santos

Reputation: 674

You can use:

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

in your example above.

Upvotes: 8

geo
geo

Reputation: 31

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.

Upvotes: 3

telliott99
telliott99

Reputation: 7907

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.

Upvotes: 2

Related Questions