Numpy：为矩阵的每列选择特定行的值

Question

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

有没有一个矢量化的方法可以做到这一点？你知道吗

Answer 1

这个怎么样：

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]: x
Out[42]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

Answer 2

作为另一种解决方案，np.choose对做出选择很有用。你知道吗

>>> x = np.arange(16).reshape(4,4)

所以x看起来像这样：

array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

现在，在每列中的特定行y处选择值可以如下所示：

>>> y = np.array([3, 0, 2, 1])
>>> np.choose(y, x)
array([12, 1, 10,  7])