如何从Python / numpy中的nd-array中的每一行中选择不同的列？

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 (eg 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 .