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

Question

假设我有3 x 2矩阵

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

并希望通过索引数组选择元素

I = [0, 1, 0]

要得到

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

我该怎么做？

这样写

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

给出了完全不同的东西（什么？）

Answer 1

你可以做的是传递一个迭代的第一dimesion值和迭代的第二个维度（如列表）。就像是：

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

这将产生：

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

如果您希望将其作为2d数组，则可以使用其他变形：

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

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

给出了完全不同的东西（什么？）

它创建一个矩阵，其中第一列是的第一（0）列A，第二列是的第二（1）列A，第三列再次是的第一（0）列A。