How can I create a vector from a matrix using spec

问题:

This question already has answers here:

using an numpy array as indices of the 2nd dim of another array? [duplicate] (2 answers)

Closed 3 years ago.

Say I have a matrix of shape (N,d) and a vector of size N which says which column in the matrix is of relevance to a given row. How can I return the vector of size N which is given by the values in the matrix and the relevant column?

For example:

M = [[ 2, 4, 1, 8],
    [3, 5, 7, 1],
    [2, 5, 3, 9],
    [1, 2, 3, 4]]
V = [2, 1, 0, 1]

I tried something like:

M[:,V]

but this returns a matrix which is NXN

Is there a simple way to format this which does not involve writing a for-loop so that I could get the following vector:

V' = [1,5,2,2]

回答1:

Use np.arange(len(V)) for indexing the row numbers and V for columns:

In [110]: M = [[ 2, 4, 1, 8],
   .....:     [3, 5, 7, 1],
   .....:     [2, 5, 3, 9],
   .....:     [1, 2, 3, 4]]

In [111]: V = [2, 1, 0, 1]

In [112]: 

In [112]: M = np.array(M)

In [113]: M[np.arange(len(V)),V]
Out[113]: array([1, 5, 2, 2])