CODEWITHSUNDEEP
pythonarraysnumpyindexing
Dims
Dims

Reputation: 51039

How to select different columns from each row in nd-array in Python/numpy?

Suppose I have 3 x 2 matrix 

A = np.arange(3*2).reshape(3,2)

and wish to select elements by index array

I = [0, 1, 0]

to get

[[0],[3],[4]]

How would I do this?

Writing this way

A[:,[0,1,0]]

gives something completely different (what?)

Upvotes: 9

Views: 3503

Answers (1)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476659

What you can do is pass an iterable of the first dimesion value, and an iterable (e.g. a list) of the second dimension. Something like:

I = [0, 1, 0]
A[range(len(I)),I]

This produces:

>>> A[range(len(I)),I]
array([0, 3, 4])

In case you want it as a 2d array, you can use an additional reshape:

>>> A[range(len(I)),I].reshape(-1,1)
array([[0],
       [3],
       [4]])

A[:,[0,1,0]]

gives something completely different (what?)

It creates a matrix where the first column is the first (0) column of A, the second column is the second (1) column of A, and the third column is again the first (0) column of A.

Upvotes: 12

Related Questions