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

Question

I have a 2D matrix X = ((a11, a12, .. a1n), (a21 .. a2n) .. (am1, .. amn)) and a 1D vector y = [y1, ..., yn] each yi is between 1 and m . For each column i of X I want to pick out the element at row yi . That is, I want to pick out the vector z = (a_(y1 1), ... a_(yn n)) .

Is there a vectorized way to do this?

Answer 1

How about this:

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

As an alternative solution, np.choose is useful for making the selections.

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

So x looks like this:

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

Now the selection of the value at a particular row y in each column can be done like this:

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