Python numpy中的高效实现（矢量化）

我在numpy有一个矩阵

A = [[1,2,3],
     [4,5,6],
     [7,8,9]]

和一个普通数组

B = [2,1,0]

我想将矩阵压缩成一列，这样对于第I行，我想选取第B[I]列，并将其放入列向量中

因此，我希望得到一个优雅的/矢量化的实现

Ans = [3,5,7]

试试这个：

In [32]: A[np.arange(A.shape[0]), B]
Out[32]: array([3, 5, 7])

---

我尝试使用索引切片而不是arange，类似于A[：，B]/A[：，B[：]的东西当然不起作用！这是优雅的谢谢@纳曼，很高兴这有帮助：）