Функция argsort

Функция argsort возвращает индексы элементов массива в том порядке, в котором они должны следовать, чтобы массив оказался отсортированным. Первым параметром функция принимает массив. Вторым параметром можно задать ось для многомерных массивов. Третьим параметром можно указать вид сортировки: 'quicksort', 'mergesort', 'heapsort' или 'stable'.

Синтаксис

np.argsort(a, [axis], [kind], [order])

Пример

Давайте получим индексы для сортировки одномерного массива:

import numpy as np

arr = np.array([3, 1, 4, 1, 5])
res = np.argsort(arr)
print(res)

Результат выполнения кода:

[1 3 0 2 4]

Полученные индексы означают, что для сортировки массива нужно взять элементы с индексами 1, 3, 0, 2, 4:

import numpy as np

arr = np.array([3, 1, 4, 1, 5])
res = np.argsort(arr)
sorted_arr = arr[res]
print(sorted_arr)

Результат выполнения кода:

[1 1 3 4 5]

Пример

Используем функцию для двумерного массива с указанием оси:

import numpy as np

arr = np.array([[3, 1, 4],
				[1, 5, 2]])
res = np.argsort(arr, axis=0)
print(res)

Результат выполнения кода:

[[1 0 1]
 [0 1 0]]

Теперь отсортируем по строкам (axis=1):

import numpy as np

arr = np.array([[3, 1, 4],
				[1, 5, 2]])
res = np.argsort(arr, axis=1)
print(res)

Результат выполнения кода:

[[1 0 2]
 [0 2 1]]

Пример

Применим разные виды сортировки для наглядности:

import numpy as np

arr = np.array([3, 1, 4, 1, 5])
res_quick = np.argsort(arr, kind='quicksort')
res_merge = np.argsort(arr, kind='mergesort')
print(res_quick)
print(res_merge)

Результат выполнения кода:

[1 3 0 2 4]
[1 3 0 2 4]

Смотрите также

функцию sort,
которая возвращает отсортированную копию массива
функцию lexsort,
которая выполняет косвенную сортировку по нескольким ключам
функцию partition,
которая выполняет частичную сортировку массива
функцию argpartition,
которая возвращает индексы для частичной сортировки