denys.fridman - 2 months ago 4x
Python Question

# mapping over 2 numpy.ndarray simultaneously

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])`
. I want to turn a 3x3 matrix A to an array of length 3 (call it v) where v[j] = A[j, p[j]] for j = 0, 1, 2. I can do it the following way:

``````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])`
(ie 0th element of row 0, 2nd element of row 1, 1st element of row 2).

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

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]
``````
Source (Stackoverflow)