resistancefm - 8 months ago 36

Python Question

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]`

Answer

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