How do I sort the rows of a 2d numpy array based on indices given by another 2d numpy array

Question

Example:

arr = np.array([[.5, .25, .19, .05, .01],[.25, .5, .19, .05, .01],[.5, .25, .19, .05, .01]])
print(arr)
[[ 0.5   0.25  0.19  0.05  0.01]
 [ 0.25  0.5   0.19  0.05  0.01]
 [ 0.5   0.25  0.19  0.05  0.01]]

idxs = np.argsort(arr)
print(idxs)
[[4 3 2 1 0]
 [4 3 2 0 1]
 [4 3 2 1 0]]

How can I use idxs to index arr ? I want to do something like arr[idxs] , but this does not work.

Answer 1

It's not the prettiest, but I think something like

>>> arr[np.arange(len(arr))[:,None], idxs]
array([[ 0.01,  0.05,  0.19,  0.25,  0.5 ],
       [ 0.01,  0.05,  0.19,  0.25,  0.5 ],
       [ 0.01,  0.05,  0.19,  0.25,  0.5 ]])

should work. The first term gives the x coordinates we want (using broadcasting over the last singleton axis):

>>> np.arange(len(arr))[:,None]
array([[0],
       [1],
       [2]])

with idxs providing the y coordinates. Note that if we had used unravel_index , the x coordinates to use would always have been 0 instead:

>>> np.unravel_index(idxs, arr.shape)[0]
array([[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]])

Answer 2

How about something like this:

I changed variables to make the example more clear, but you basically need to index by two 2D arrays.

In [102]: a = np.array([[1,2,3], [4,5,6]])

In [103]: b = np.array([[0,2,1], [2,1,0]])

In [104]: temp = np.repeat(np.arange(a.shape[0]), a.shape[1]).reshape(a.shape).T 

          # temp is just [[0,1], [0,1], [0,1]]
          # probably can be done more elegantly

In [105]: a[temp, b.T].T
Out[105]: 
array([[1, 3, 2],
       [6, 5, 4]])