2D Numpy 数组的高级索引

Question

我有一个大小为 [N, M] = A 的 2D numpy 数组。 我试图根据某些条件从每一行中选择一个元素。 我已经拥有大小为 [N, 1] = B 的数组中的所有索引，其中该数组中的每个元素都在 [0, M) 中。 即B中的每个数字表示每行中待选择元素的位置。 我已经尝试过 A[B]，但没有达到预期效果。

Answer 1

您可以使用

np.arrange(N)

生成行索引序列，然后将它们与

B

一起使用来正确索引

A

：

N, M = A.shape
indices = np.arange(N)
selected = A[indices, B.flatten()]