Сортировки в Python

В Python сортировка устроена иначе, чем в C++, и отличия не только в синтаксисе. Кое-что здесь гарантировано стандартом, чего в C++ нет.

sorted и list.sort

b = sorted(a)      # новый список, исходный не меняется
a.sort()           # на месте, возвращает None

Частая ошибка — написать a = a.sort() и получить None. Метод сортирует на месте и ничего не возвращает намеренно: так видно, что объект изменён.

sorted работает с любым перебираемым объектом и всегда отдаёт список: sorted("сортировка") вернёт список букв, а не строку.

Timsort

Внутри — Timsort, и у него три свойства, которые стоит знать.

Устойчивость гарантирована. Не «так вышло в этой реализации», а записано в документации языка. В C++ такой гарантии у sort нет — там для этого есть отдельный stable_sort.

$O(n \log n)$ в худшем случае гарантированы. Развернуть Python-сортировку в квадрат подобранным входом нельзя — в отличие от наивной быстрой сортировки.

Готовые отсортированные куски используются. Timsort ищет в массиве уже упорядоченные участки и сливает их, вместо того чтобы сортировать заново. На почти отсортированных данных он работает почти за линию — то же свойство, что у сортировки вставками, только в промышленном исполнении.

Ключи

Порядок задаётся не компаратором, а ключом — функцией, которая по элементу выдаёт то, по чему сравнивать.

people.sort(key=lambda p: p.age)              # по возрасту
people.sort(key=lambda p: (-p.score, p.name)) # по убыванию баллов, потом по имени
words.sort(key=len)                           # по длине
a.sort(reverse=True)                          # по убыванию

Ключ вычисляется по одному разу на элемент, а не на каждое сравнение. Это важно, когда ключ дорогой: key=lambda s: sorted(s) для группировки анаграмм посчитается $n$ раз, а не $n \log n$.

Сортировка по нескольким полям делается кортежем. Для чисел убывание задаётся минусом, для строк — нет; там либо reverse=True на весь порядок, либо два прохода с опорой на устойчивость:

people.sort(key=lambda p: p.name)             # сначала младший ключ
people.sort(key=lambda p: p.score, reverse=True)  # потом старший

Второй проход не разрушит порядок, наведённый первым, — именно потому, что сортировка устойчива.

Когда ключа не хватает

Иногда порядок нельзя выразить ключом. Классический случай — склеить куски в наибольшее число: там нужно сравнение «$x + y$ против $y + x$», а не значение элемента.

import functools

pieces.sort(key=functools.cmp_to_key(
    lambda x, y: -1 if x + y > y + x else (1 if x + y < y + x else 0)
))

cmp_to_key оборачивает компаратор в объект, у которого определён <. Работает, но заметно медленнее ключа: сравнение вызывается $O(n \log n)$ раз, и каждый раз это вызов Python-функции.

В отличие от C++, нарушение строгости здесь программу не роняет — просто порядок получится не тот, какой вы хотели.

heapq

Модуль heapq работает с обычным списком как с min-кучей.

import heapq

heapq.heapify(a)            # превратить список в кучу на месте, O(n)
heapq.heappush(a, x)
smallest = heapq.heappop(a)

heapq.nsmallest(10, a)      # десять наименьших
heapq.nlargest(10, a)       # десять наибольших

nsmallest(k, a) работает за $O(n \log k)$ — держит кучу из $k$ элементов и не сортирует весь массив. При маленьком $k$ это заметно быстрее, чем sorted(a)[:k].

Max-кучи в модуле нет. Обходятся хранением отрицательных значений или кортежей (-priority, item).

bisect

Двоичный поиск по отсортированному списку:

from bisect import bisect_left, bisect_right, insort

i = bisect_left(a, x)       # первая позиция, куда можно вставить x
j = bisect_right(a, x)      # последняя такая позиция
count = j - i               # сколько раз x встречается

Есть и insort, вставляющий элемент с сохранением порядка. Но осторожно: сам поиск позиции логарифмический, а вот вставка в список сдвигает хвост и стоит $O(n)$. Для поддержания отсортированного набора из многих элементов это плохой выбор.

Чего в Python нет

Прямого аналога nth_element в стандартной библиотеке не существует. Если нужен ровно $k$-й элемент, ближайшее — heapq.nsmallest(k, a)[-1] за $O(n \log k)$ или собственный quickselect из статьи про порядковые статистики.

И общее замечание про скорость: сортировка в Python выполняется скомпилированным кодом, а вот ключ — обычная Python-функция, и вызывается она $n$ раз. Если сортировка тормозит, дело почти всегда в ключе, а не в самой сортировке.