我有一个二维矩阵 X = ((a11, a12, .. a1n), (a21 .. a2n) .. (am1, .. amn)) 和一个一维向量 y = [y1, ..., yn]，每个 yi 都在 1 和 m 之间。对于 X 的每一列 i，我想选出第 yi 行的元素。也就是说，我想选出向量 z = (a_(y1 1), ... a_(yn n))。

有没有一种向量化的方法来实现这个操作？

回答：

这样做如何：

In [39]: x = np.arange(12).reshape(4,3)In [40]: y = np.array([0,3,2])In [41]: x[y[None, :], np.arange(len(y))[None,:]][0]Out[41]: array([ 0, 10,  8])In [42]: xOut[42]: array([[ 0,  1,  2],       [ 3,  4,  5],       [ 6,  7,  8],       [ 9, 10, 11]])