numpy.argsort

numpy.argsort(a, axis=-1, kind='quicksort', order=None)

Функция argsort() возвращает индексы, сортирующие элементы исходного массива.

Параметры:
a - массив NumPy или подобный массиву объект.
Исходный массив.
axis - целое число (необязательный параметр).
Определяет ось вдоль которой выполняется сортировка элементов. Если равен None, то сортировка выполняется по сжатому до одной оси представлению исходного массива. По умолчанию axis = -1, что соответствует сортировке по последней оси массива.
kind - строка (необязательный параметр).
Одна строка из множества {'quicksort', 'mergesort', 'heapsort', 'stable'}, которая указывает на алгоритм сортировки. Значение 'stable' приводит к автоматическому выбору устойчивого, лучшего алгоритма для типа данных сортируемого массива. По умолчанию kind = 'quicksort'.
order - строка или список строк (необязательный параметр).
В случае, когда массив a является структурированным, и его полям присвоено имя, то данный параметр позволяет указать порядок в котором они должны учавствовать в сортировке Если указаны не все поля, то неуказанные будут использоваться в том порядке в котором они перечислены в dtype структурированного массива.
Возвращает:
ndarray - массив NumPy
Массив индексов, которые сортируют элементы исходного массива вдоль указанной оси axis. Если a это одномерный массив, то a[np.argsort(a)] вернет его отсортированным. Получить отсортированные многомерные массивы можно выполнив либо итерирование по ним с использованием возвращаемого массива индексов, либо воспользовавшись командой np.take_along_axis(a, np.argsort(a, axis = ax), axis=ax).
Смотрите так же:
sort, lexsort, argpartition


Примеры

Одномерные массивы:

>>> import numpy as np
>>> 
>>> a = np.random.randint(0, 20, 10)
>>> a
array([ 8, 17,  8, 14, 14, 16,  4,  3, 12,  6])
>>> 
>>> ind = np.argsort(a)
>>> ind
array([7, 6, 9, 0, 2, 8, 3, 4, 5, 1], dtype=int32)
>>> 
>>> a[ind]
array([ 3,  4,  6,  8,  8, 12, 14, 14, 16, 17])

Двумерные массивы:

>>> a = np.random.randint(0, 20, (3, 4))
>>> a
array([[19,  3,  5,  6],
       [ 0,  0,  3,  2],
       [13, 14,  8, 13]])
>>> 
>>>               
>>> #  Индексы отсортированных строк:
... ind_r = np.argsort(a)    #  эквивалентно ind_r = np.argsort(a, axis = 1)
>>> ind_r
array([[1, 2, 3, 0],
       [0, 1, 3, 2],
       [2, 0, 3, 1]], dtype=int32)
>>> 
>>> 
>>> #  Индексы отсортированных столбцов:
... ind_c = np.argsort(a, axis = 0)
>>> ind_c
array([[1, 1, 1, 1],
       [2, 0, 0, 0],
       [0, 2, 2, 2]], dtype=int32)

Самый простой способ получить отсортированный массив - воспользоваться функцией np.take_along_axis():

>>> np.take_along_axis(a, ind_r, axis = 1)
array([[ 3,  5,  6, 19],
       [ 0,  0,  2,  3],
       [ 8, 13, 13, 14]])
>>> 
>>> np.take_along_axis(a, ind_c, axis = 0)
array([[ 0,  0,  3,  2],
       [13,  3,  5,  6],
       [19, 14,  8, 13]])

Функция np.argsort() так же позволяет выполнить лексикографическую сортировку структурированного массива:

>>> arr = np.array([(0, 12), (0, 11), (1, 55),
...                 (0, 87), (1, 54), (1, 21)],
...                dtype = np.dtype([('a', int), ('b', int)]))
>>> arr
array([(0, 12), (0, 11), (1, 55), (0, 87), (1, 54), (1, 21)],
      dtype=[('a', '<i4'), ('b', '<i4')])
>>> 
>>> ind = np.argsort(arr, order = ['a', 'b'])
>>> ind
array([1, 0, 3, 5, 4, 2], dtype=int32)
>>> 
>>> arr[ind]
array([(0, 11), (0, 12), (0, 87), (1, 21), (1, 54), (1, 55)],
      dtype=[('a', '<i4'), ('b', '<i4')])
>>> 
>>> 
>>> ind = np.argsort(arr, order = ['b', 'a'])
>>> ind
array([1, 0, 5, 4, 2, 3], dtype=int32)
>>> 
>>> arr[ind]
array([(0, 11), (0, 12), (1, 21), (1, 54), (1, 55), (0, 87)],
      dtype=[('a', '<i4'), ('b', '<i4')])