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]]
@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.]])
Not pretty but:
A[np.arange(2),B.T].T
array([[ 25., 13.],
[ 18., 25.]])
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.