Question

我有一个2D矩阵X = ((a11, a12, .. a1n), (a21 .. a2n) .. (am1, .. amn))和一个y = [y1, ..., yn] {1}} yi介于1和m之间。对于i的每个X列，我想在行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])

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

2 个答案: