Have numpy argsort return an array of 2d indices?

Question

If we have a 1d array

arr = np.random.randint(7, size=(5)) # [3 1 4 6 2] print np.argsort(arr) # [1 4 0 2 3] <= The indices in the sorted order     

If we have a 2d array

arr = np.random.randint(7, size=(3, 3)) # [[5 2 4] # [3 3 3] # [6 1 2]] print np.argsort(arr) # [[1 2 0] # [0 1 2] # [1 2 0]] <= It sorts each row 

What I need is the 2d indices that sort this matrix in its entirety. Something like this:

# [[2 1] => 1 # [0 1] => 2 # [2 2] => 2 # . # . # . # [0 2] => 4 # [0 0] => 5 # [2 0]] => 6 

How do I get "2d indices" for the sorting of a 2d array?

Answer 1

Apply numpy.argsort on flattened array and then unravel the indices back to (3, 3) shape:

>>> arr = np.array([[5, 2, 4], [3, 3, 3], [6, 1, 2]]) >>> np.dstack(np.unravel_index(np.argsort(arr.ravel()), (3, 3))) array([[[2, 1],         [0, 1],         [2, 2],         [1, 0],         [1, 1],         [1, 2],         [0, 2],         [0, 0],         [2, 0]]])