python怎么选取不连续的列,如何从Python / numpy中的nd-array中的每一行中选择不同的列？...

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

解决方案

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.