2D 到 1D numpy 数组，每行的列索引

Question

我正在尝试为矩阵中的每一行索引一列。

假设我有一个形状为(n,m)的 numpy 数组A 。
我还有一个 numpy 数组B ，其形状(n,)包含 0 到 m 之间的整数，因此它们可以用作A的第二轴的索引。

我想获得一个 numpy 数组C ，其形状为(n,) ， C[i] = A[i,B[i]] ，因此A每一行都会根据B中的索引产生一个值。

当然，我可以在 for 循环或列表理解中使用最后一个表达式，但是我将如何使用 numpy 的效率来做到这一点？

我的第一个直觉是C = A[:,B]但这显然会产生其他东西。 （形状(n,n) ）

Answer 1

你很接近，试试：

C = A[range(len(B)), B]

这应该产生C[i] = A[i,B[i]]