Numpy:select 矩阵每一列特定行的值

Numpy: select value at a particular row for each column of a matrix

我有一个二维矩阵 X = ((a11, a12, .. a1n), (a21 .. a2n) .. (am1, .. amn)) 和一个一维向量 y = [y1, ..., yn] 每个 yi 都在 1m 之间。对于 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]: x
Out[42]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

作为替代解决方案,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])