Numpy: Sorting a multidimensional array by a multidimensional array
- by JD Long
Forgive me if this is redundant or super basic. I'm coming to Python/Numpy from R and having a hard time flipping things around in my head.
I have a n dimensional array which I want to sort using another n dimensional array of index values. I know I could wrap this in a loop but it seems like there should be a really concise Numpyonic way of beating this into submission. Here's my example code to set up the problem where n=2:
a1 = random.standard_normal(size=[2,5])
index = array([[0,1,2,4,3] , [0,1,2,3,4] ])
so now I have a 2 x 5 array of random numbers and a 2 x 5 index. I've read the help for take() about 10 times now but my brain is not groking it, obviously.
I thought this might get me there:
take(a1, index)
array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984, 0.25698875],
[ 0.29589188, -0.71279375, -0.18154864, 0.25698875, -1.12184984]])
but that's clearly reordering only the first element (I presume because of flattening).
Any tips on how I get from where I am to a solution that sorts element 0 of a1 by element 0 of the index ... element n?