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

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

您可以做的是传递第一个维度值的 iterable 和第二个维度的 iterable(例如列表)。像这样的东西:

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

这会产生:

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

如果你想要它作为一个二维数组，你可以使用额外的 reshape :

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

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

gives something completely different (what?)

它创建一个矩阵，其中 first 列是 A 的第一 (0) 列，第二> 列是 A 的第二 (1) 列，第三列也是 A 的第一 (0) 列。