How Can I Create A Vector From A Matrix Using Specified Colums From Each Row Without Looping In Python?

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

Solution 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])