Нужно ли numpy argsort возвращать массив из двух индексов?

Question

Если мы имеем 1D массива

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  

Если у нас есть 2D-массив

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 

Что мне нужно это 2d индексы, сортировать эту матрица в полном комплекте. Что-то вроде этого:

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

Как получить «2d-индексы» для сортировки 2-мерного массива?

Answer 1

Применить numpy.argsort на плоский массив, а затем распутать индексы обратно (3, 3) формы:

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