Row-wise indexing in Numpy

Question

I have two matrices, A and B:

A = array([[2., 13., 25., 1.], [ 18., 5., 1., 25.]])
B = array([[2, 1], [0, 3]])

I want to index each row of A with each row of B, producing the slice:

array([[25., 13.], [18., 25.]])

That is, I essentially want something like:

array([A[i,b] for i,b in enumerate(B)])

Is there a way to fancy-index this directly? The best I can do is this "flat-hack":

A.flat[B + arange(0,A.size,A.shape[1])[:,None]]

Answer 1

@Ophion's answer is great, and deserves the credit, but I wanted to add some explanation, and offer a more intuitive construction.

Instead of rotating B and then rotating the result back, it's better to just rotate the arange. I think this gives the most intuitive solution, even if it takes more characters:

A[((0,),(1,)), B]

or equivalently

A[np.arange(2)[:, None], B]

This works because what's really going on here, is you're making an i array and a j array, each of which have the same shape as your desired result.

i = np.array([[0, 0],
              [1, 1]])
j = B

But you can use just

i = np.array([[0],
              [1]])

Because it will broadcast to match B (this is what np.arange(2)[:,None] gives).

Finally, to make it more general (not knowing 2 as the arange size), you could also generate i from B with

i = np.indices(B.shape)[0]

however you build i and j, you just call it like

>>> A[i, j]
array([[ 25.,  13.],
       [ 18.,  25.]])