基于行索引从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部分开始：

  

当索引由整数数组组成时   被索引有维度，索引很简单，但是   与切片不同。 [...]最好用一个例子来理解。

Answer 2

索引2D numpy数组可能有点令人困惑。

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