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

Question

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?)

Answer 1

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.