基于行索引的numpy数组列元素选择

Question

我有一个二维数组A和一个索引列表idx，例如：

A = np.array([[ 1.,  0.,  0.],
             [ 0.,  1.,  0.],
             [ 0.,  0.,  1.],
             [ 0., -1.,  0.],
             [ 0.,  0.,  5.]])

idx = np.array([2, 1, 0, 1, 2])

我试图选择A的元素，这些元素沿着列轴被idx索引（在这个例子中：array([0., 1., 0., -1., 5.])）。我怎么能做到这一点没有循环？你知道吗

谢谢你！你知道吗

Answer 1

A[np.arange(np.size(idx)), idx]

给出array([ 0., 1., 0., -1., 5.])

从文档的Advanced Indexing部分：

When the index consists of as many integer arrays as the array being indexed has dimensions, the indexing is straight forward, but different from slicing. [...] This is best understood with an example.

Answer 2

索引2D numpy数组可能有点混乱。你知道吗

你需要：A[np.arange(0, A.shape[0]), idx]