Сортировка и поиск: sort, argsort, searchsorted, unique

Урок про упорядочивание и поиск: как сортировать, получать порядок индексов и эффективно искать в отсортированных данных.

argsort возвращает не сами отсортированные значения, а индексы, в порядке которых элементы образуют отсортированную последовательность; это ключ к сортировке связанных массивов.

Сортировка значений: sort

Сортировка и поиск — базовые операции, без которых не обходится анализ данных: ранжирование, поиск порогов, дедупликация, подсчёт частот. NumPy реализует их векторно и быстро, так что упорядочить миллионы чисел или найти позицию в отсортированном массиве — дело одного вызова, а не цикла. np.sort(a) возвращает новый отсортированный массив (копию), не меняя исходный. Метод a.sort() сортирует на месте. Для многомерных массивов сортировка по умолчанию идёт вдоль последней оси, но ось можно задать через axis.

import numpy as np
a = np.array([3, 1, 4, 1, 5, 9, 2, 6])

print(np.sort(a))        # новый отсортированный массив
print(a)                 # оригинал не тронут

m = np.array([[3, 1, 2],
              [9, 5, 7]])
print(np.sort(m, axis=1))  # сортировка внутри каждой строки

Вывод:

[1 1 2 3 4 5 6 9]
[3 1 4 1 5 9 2 6]
[[1 2 3]
 [5 7 9]]

argsort: порядок индексов

Это, пожалуй, самая важная идея урока, поэтому остановимся на ней основательно. Часто нужно не просто отсортировать массив, а узнать порядок, в котором его элементы идут по возрастанию — и применить этот порядок к другому связанному массиву. Например, отсортировать имена по их оценкам. argsort возвращает индексы, которые расставляют элементы по порядку.

import numpy as np
scores = np.array([50, 90, 70, 60])
names = np.array(['Аня', 'Боря', 'Вася', 'Галя'])

order = np.argsort(scores)      # индексы для сортировки по возрастанию
print(order)                    # [0 3 2 1]
print(scores[order])            # отсортированные оценки
print(names[order])             # имена в том же порядке!

Вывод:

[0 3 2 1]
[50 60 70 90]
['Аня' 'Галя' 'Вася' 'Боря']

Это мощнейший приём: argsort даёт «перестановку», которую можно применить к любому параллельному массиву. Для сортировки по убыванию используют scores[np.argsort(scores)[::-1]]. Воспроизведём идею argsort на чистом Python — сортируем индексы по ключу-значению:

scores = [50, 90, 70, 60]
names = ['Аня', 'Боря', 'Вася', 'Галя']

order = sorted(range(len(scores)), key=lambda i: scores[i])
print("Порядок индексов:", order)
print("По оценкам:", [scores[i] for i in order])
print("Имена:    ", [names[i] for i in order])

Вывод:

Порядок индексов: [0, 3, 2, 1]
По оценкам: [50, 60, 70, 90]
Имена:     ['Аня', 'Галя', 'Вася', 'Боря']

Стабильность сортировки и выбор алгоритма

У сортировки в NumPy есть параметр kind, задающий алгоритм: quicksort (по умолчанию, быстрый в среднем), mergesort (стабильный), heapsort, stable. Большинству задач хватает значения по умолчанию, но в одном случае выбор важен — когда нужна стабильная сортировка. Стабильная сортировка сохраняет относительный порядок элементов с равными ключами. Это критично при многоуровневой сортировке: чтобы отсортировать данные сначала по одному полю, потом по другому, вы сортируете сначала по менее важному ключу, затем стабильно по более важному — и порядок по первому ключу сохранится внутри групп с равным вторым. Для этого берут kind='stable' (или mergesort). Знание о стабильности также объясняет, почему argsort с равными значениями возвращает индексы в определённом порядке. Для разовой сортировки одного массива это неважно, но для сложных многокритериальных упорядочиваний стабильность — необходимое свойство, и хорошо помнить, как её включить.

Поиск в отсортированном массиве: searchsorted

Если массив уже отсортирован, искать в нём можно за логарифмическое время бинарным поиском. np.searchsorted(a, v) возвращает индекс, куда нужно вставить значение v, чтобы сохранить порядок. Это используют для быстрого поиска позиции, бинирования (распределения по интервалам), интерполяции.

import numpy as np
a = np.array([1, 3, 5, 7, 9])   # отсортирован

print(np.searchsorted(a, 4))        # куда вставить 4 -> индекс 2
print(np.searchsorted(a, 7))        # позиция существующего -> 3
print(np.searchsorted(a, [0, 6, 10]))  # для нескольких значений сразу

Вывод:

2
3
[0 3 5]

Бинарный поиск работает только на отсортированных данных, но зато очень быстр — логарифм против линейного перебора. Покажем его суть на Python:

def searchsorted(a, v):
    lo, hi = 0, len(a)
    while lo < hi:
        mid = (lo + hi) // 2
        if a[mid] < v:
            lo = mid + 1
        else:
            hi = mid
    return lo

a = [1, 3, 5, 7, 9]
print(searchsorted(a, 4))   # 2
print(searchsorted(a, 6))   # 3
print(searchsorted(a, 10))  # 5

Вывод:

2
3
5

Почему argsort — это «суперсила»

Стоит задержаться на argsort, потому что это один из самых недооценённых инструментов NumPy. Его сила в том, что он отделяет порядок от данных. Получив перестановку индексов, вы можете применить её к чему угодно: к самим значениям, к параллельным массивам имён, дат, меток, к нескольким столбцам сразу. Это решает целый класс задач «отсортировать одно по другому»: расставить товары по продажам, студентов по баллам, события по времени, сохранив связь всех сопутствующих данных. Более того, через argsort выражаются и более хитрые операции: топ-k элементов (взять последние k индексов отсортированного порядка), ранги элементов (argsort от argsort даёт позицию каждого элемента в отсортированном порядке), сортировка по нескольким ключам. Освоив мышление «сначала получаю порядок индексов, потом применяю его», вы решаете задачи упорядочивания элегантно и без циклов. Это тот приём, который отличает беглый код на NumPy: вместо того чтобы тасовать данные напрямую, вы вычисляете перестановку и применяете её ко всему, что нужно синхронизировать.

Уникальные значения: unique

np.unique возвращает отсортированные уникальные значения массива. С дополнительными параметрами он также считает, сколько раз встречается каждое значение (return_counts=True), и даёт обратные индексы для восстановления. Это незаменимо для подсчёта частот, поиска категорий, удаления дубликатов.

import numpy as np
a = np.array([3, 1, 2, 3, 3, 1, 2, 2, 2])

print(np.unique(a))                          # [1 2 3]
values, counts = np.unique(a, return_counts=True)
print(values)
print(counts)                                # сколько раз каждое
# Самое частое значение:
print(values[counts.argmax()])

Вывод:

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

Связка unique(return_counts=True) + argmax — классический способ найти моду (самое частое значение). Аналог на чистом Python через Counter:

from collections import Counter

a = [3, 1, 2, 3, 3, 1, 2, 2, 2]
c = Counter(a)
print("Частоты:", dict(sorted(c.items())))
print("Самое частое:", c.most_common(1)[0][0])

Вывод:

Частоты: {1: 2, 2: 4, 3: 3}
Самое частое: 2

unique как инструмент анализа категорий

np.unique — гораздо больше, чем «удалить дубликаты». Это рабочая лошадка анализа категориальных данных. С return_counts=True он строит частотную таблицу: какие значения встречаются и сколько раз — основа для гистограмм категорий, поиска самых частых элементов, обнаружения редких. С return_index=True он сообщает, где впервые встретилось каждое уникальное значение. С return_inverse=True он даёт массив, которым можно восстановить исходные данные из уникальных, — это, по сути, кодирование категорий в целые числа, шаг, который постоянно нужен при подготовке данных для моделей (превратить названия категорий в номера). Таким образом, одна функция покрывает целый спектр задач: дедупликацию, подсчёт частот, кодирование категорий, поиск моды. Когда вам нужно понять «какие вообще значения есть в этих данных и в каком количестве», unique — первый инструмент, к которому стоит обратиться. И снова это векторная операция: она обрабатывает миллионы элементов одним вызовом, без циклов и словарей.

Частичная сортировка: partition

Иногда полная сортировка избыточна — нужны, скажем, только k наименьших элементов, а их внутренний порядок неважен. np.partition делает это быстрее полной сортировки: он расставляет массив так, что k-й элемент стоит на своём итоговом месте, слева — меньшие, справа — большие (но не отсортированные). Полезно для поиска топ-k.

import numpy as np
a = np.array([7, 2, 9, 1, 5, 3, 8])

# 3 наименьших элемента окажутся слева (в произвольном порядке)
print(np.partition(a, 3))
print(np.sort(np.partition(a, 3)[:3]))   # сами 3 наименьших

Вывод:

[1 2 3 5 7 9 8]
[1 2 3]

Подводные камни

• Путать sort и argsort. Первый даёт значения, второй — индексы. Для связанных массивов нужен именно argsort.
• searchsorted на неотсортированном. Бинарный поиск требует отсортированного массива; иначе результат бессмыслен.
• a.sort() меняет оригинал. Метод сортирует на месте; функция np.sort возвращает копию.
• Ось сортировки. По умолчанию сортируется последняя ось; для другой задавайте axis явно.

searchsorted для бинирования и быстрого поиска

Помимо очевидного «найти позицию», searchsorted решает важную задачу бинирования — распределения значений по интервалам. Если у вас есть отсортированные границы корзин (например, возрастные группы 0–18, 18–30, 30–50, 50+), то np.searchsorted(границы, значения) сразу вернёт для каждого значения номер корзины, в которую оно попадает. Это векторная замена цепочке условий и работает за логарифмическое время на каждый элемент. Тот же приём лежит в основе построения гистограмм и кусочно-линейной интерполяции: найти, между какими узлами сетки лежит точка. Ключевое требование — границы должны быть отсортированы, иначе бинарный поиск даст бессмыслицу. Поскольку searchsorted принимает сразу массив искомых значений, тысячи и миллионы точек бинируются одним вызовом без цикла. Это один из тех инструментов, которые превращают типичную «циклическую» задачу (для каждого значения перебрать интервалы) в чистую векторную операцию — и стоит держать его в арсенале для любых задач классификации по диапазонам.

Лучшие практики

• Для сортировки одного массива по ключу из другого используйте argsort и применяйте порядок индексов.
• В отсортированных данных ищите через searchsorted — это логарифмически быстро.
• Частоты и категории считайте через np.unique(return_counts=True).
• Для топ-k берите np.partition вместо полной сортировки.

Сортировка и поиск — это инструменты упорядочивания и навигации по данным. Главные идеи урока: отделяйте порядок от данных через argsort, используйте логарифмически быстрый searchsorted на отсортированных массивах, анализируйте категории через unique. Все эти операции векторные и работают на больших данных без циклов, что делает их незаменимыми в реальном анализе.

Итог

• np.sort возвращает отсортированную копию, a.sort() — сортирует на месте.
• argsort даёт порядок индексов — ключ к сортировке связанных массивов.
• searchsorted — бинарный поиск позиции в отсортированном массиве.
• unique(return_counts=True) считает частоты; partition быстро находит топ-k.