Czy numpy argsort zwraca tablicę indeksów 2d?

Question

Jeśli mamy 1D tablicę

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  

Jeśli mamy tablicy 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 

Co potrzebne jest 2D indeksy że porządek ten matryca w jego całości. Coś takiego:

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

Jak uzyskać "wskaźniki 2d" do sortowania tablicy 2d?

Answer 1

Zastosuj numpy.argsort na spłaszczonym tablicy i następnie rozwikłać indeksy z powrotem do (3, 3) kształt:

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