Python numpy сохраняет список индексов отсортированного 2D-массива

у меня есть массив 2D numpy, и я хочу создать новый массив 1D, где это индексы чисел в первом массиве, если они отсортированы в порядке возрастания. Для следующего массива:

A = [[1,0,2],
     [0,3,0]]

Я хочу, чтобы это было так:

B = [[1,1],[0,2],[0,0],[0,1],[1,0],[1,2]]

любая идея, как это можно сделать в python, используя предопределенные функции?

спасибо

6
arrays numpy python sorting

1 ответов

можно использовать argsort для сортировки индексов сплющенного массива, а затем unravel_index чтобы преобразовать плоский индекс обратно в координаты:

>>> i = (-a).argsort(axis=None, kind='mergesort')
>>> j = np.unravel_index(i, a.shape) 
>>> np.vstack(j).T
array([[1, 1],
       [0, 2],
       [0, 0],
       [0, 1],
       [1, 0],
       [1, 2]])

-a и kind='mergesort' - это для того, чтобы отсортировать массив в стабильный манера в порядке убывания (чтобы соответствовать выходу, который вы ищете).

если вы не заботитесь о стабильной сортировке, замените первую строку на:

>>> i = a.argsort(axis=None)[::-1]
7