denys.fridman denys.fridman - 4 months ago 12
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.

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]