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    
         


        

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]]])