denys.fridman - 4 months ago 12

Python Question

Here's the problem. Let's say I have a matrix A =

`array([[ 1., 0., 2.],`

[ 0., 0., 2.],

[ 0., -1., 3.]])

and a vector of indices p =

`array([0, 2, 1])`

`v = map(lambda (row, idx): row[idx], zip(A, p))`

So for the above matrix A and a vector of indices p I expect to get

`array([1, 2, -1])`

But can I achieve the same result by using native numpy (ie without explicitly zipping and then mapping)? Thanks.

Answer

I don't think that such a functionality exists. To achieve what you want, I can think of two easy ways. You could do:

```
np.diag(A[:, p])
```

Here the array `p`

is applied as a column index for every row such that on the diagonal you will have the elements that you are looking for.

As an alternative you can avoid to produce a lot of unnecessary entries, you can use:

```
A[np.arange(A.shape[0]), p]
```