CODEWITHSUNDEEP
pythonarraysnumpymatrix
denys.fridman
denys.fridman

Reputation: 165

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.

Upvotes: 0

Views: 75

Answers (1)

jotasi
jotasi

Reputation: 5177

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 by using:

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

Upvotes: 3

Related Questions