NumPy: вычисления и устройство массивов

Сортировка массивов: sort и argsort

Содержание курса

np.argsort: индексы порядка элементов

np.sort отвечает на вопрос «каковы значения в отсортированном порядке?». np.argsort отвечает на другой вопрос: «в каком порядке нужно взять элементы исходного массива, чтобы они оказались отсортированы?». Результат — не значения, а целочисленные индексы.

Для одномерного массива:

import numpy as np

arr = np.array([30, 10, 40, 20])
print(np.argsort(arr))  # [1 3 0 2]

Читается так: на первом месте стоит элемент с индексом 1 (это 10), на втором — с индексом 3 (это 20), на третьем — с индексом 0 (это 30), на четвёртом — с индексом 2 (это 40). Если применить этот результат обратно к arr, получим отсортированный массив:

idx = np.argsort(arr)
print(arr[idx])  # [10 20 30 40]

То есть np.argsort — это не альтернативный способ получить те же [10, 20, 30, 40], а способ узнать откуда они пришли.

Для двумерного массива поведение управляется параметром axis, как и у np.sort.

M = np.array([[3, 1, 4],
              [1, 5, 2]])

print(np.argsort(M, axis=1))
# [[1 0 2]
#  [0 2 1]]

При axis=1 каждая строка обрабатывается независимо. В строке 0 значения [3, 1, 4]: наименьшее — на позиции 1 (это 1), затем позиция 0 (это 3), затем позиция 2 (это 4). В строке 1 значения [1, 5, 2]: порядок — 0, 2, 1. Каждый индекс — это позиция внутри своей строки.

При axis=0 картина другая: каждый столбец сортируется независимо, а индексы показывают номер строки внутри своего столбца.

print(np.argsort(M, axis=0))
# [[1 0 1]
#  [0 1 0]]

Столбец 0: значения [3, 1] — наименьшее в строке 1, большее в строке 0. Столбец 1: [1, 5] — наименьшее в строке 0. Столбец 2: [4, 2] — наименьшее в строке 1.

Главное отличие от np.sort: результат np.argsort сам по себе нельзя интерпретировать как данные — это карта перестановки, инструмент для последующей индексации.