Top-k: k лучших

Какие десять землетрясений силой от 5 за 2015–2025 годы случились ближе всего к Алматы? В каталоге USGS 19 073 события. Можно посчитать расстояние до каждого и отсортировать все. А можно прочитать каталог один раз, держа в памяти только десять лучших на текущий момент: новое событие сравнивается с самым дальним из этих десяти и либо выталкивает его, либо проходит мимо. Нажмите на карту в другом месте — и посмотрите, как сжимается круг.

Из 19 073 землетрясений в кучу из десяти вошли 122, остальные 18 951 отсеяны одним сравнением с вершиной. Ближайшие — в 248–266 км от площади Республики, и шесть из десяти случились в один день, 22 января 2024 года: это землетрясение силой 7,0, которое тоже есть в списке, и толчки рядом с ним. Посмотрите на график под картой: число входов в кучу растёт всё медленнее — так, как и должно расти, по логарифму. Почему именно так, расскажем ниже.

Идея

Чтобы найти k наименьших, держим кучу из k элементов, у которой наверху наибольший из них — худший из лучших, порог входа. Очередной элемент сравниваем с порогом: если не лучше — пропускаем; если лучше — он заменяет вершину и просеивается вниз. В конце в куче k лучших, и сортируем только их.

Здесь работает та же двоичная куча, что в календарях событий, но вывернутая: для k ближайших наверху самый дальний, для k сильнейших — самый слабый. Куче не нужен весь массив сразу, ей нужен только поток. Поэтому k лучших можно искать в журнале, который не помещается в память, в данных, которые ещё приходят, и в генераторе без конца. Памяти — ровно на k элементов, сколько бы их ни было всего.

Цена — O(log k) на элемент, который вошёл в кучу, и одно сравнение на элемент, который не вошёл. В худшем случае входят все: если поток идёт от худших к лучшим, каждый новый элемент лучше всех предыдущих, и выходит n log₂ k сравнений. Но обычно входит мало.

Сколько раз меняется куча

Для n = 19 073 и k = 10 это около 85 входов, для k = 100 — около 625, для k = 1 — около 10. На каталоге вышло 122, 565 и 5: события идут не в случайном порядке по расстоянию, а толпами — сильное землетрясение и его повторные толчки лежат рядом и во времени, и на карте. Но порядок величины тот же. Значит, работы почти ровно n сравнений с вершиной плюс небольшая добавка на входы, и при малом k куча обходит сортировку почти в log₂ n раз.

Четыре способа

Кучей из k дело не исчерпывается. Если все данные уже в памяти, k лучших можно получить ещё тремя способами: отсортировать всё и взять первые k; найти k-й по величине быстрым выбором — он заодно отделит k лучших от остальных — и отсортировать только их; построить кучу из всех n за линейное время и k раз достать вершину, то есть остановить пирамидальную сортировку после k шагов. Ниже — сравнения всех четырёх на 19 073 землетрясениях, для каждого k.

При k = 10 куче нужно 19 440 сравнений — на 368 больше, чем n − 1, необходимых, чтобы найти один максимум. Сортировке — 233 244, в 12 раз больше. Быстрый выбор и куча из всех стоят около двух n при любом малом k: оба обязаны хотя бы раз переложить весь массив. С ростом k куча из k дорожает: при k = 1000 быстрый выбор уже в полтора раза дешевле её, а при k = 10 000, больше половины массива, куча из k проигрывает даже полной сортировке. Здесь работает её худший случай: когда k сравнимо с n, в кучу входит большая часть потока, и каждый вход стоит логарифм.

Код

Десять ближайших к Алматы — первая сцена в виде программы. В Python куча — модуль heapq, он умеет только кучу с наименьшим наверху, поэтому расстояние кладётся с минусом. В Go — container/heap и функция heap.Fix: новый элемент записывается прямо в вершину и просеивается вниз, одной операцией вместо двух.

Все четыре способа со счётчиком сравнений — те же числа, что во второй сцене. Проверка в конце следит, чтобы все четыре дали один и тот же ответ.

Сколько это стоит

Куча из k: время от O(n) при удачном порядке до O(n log k) в худшем случае, в среднем на случайном порядке — n + O(k log k · log(n/k)) сравнений; память O(k). Быстрый выбор с сортировкой лучших — O(n + k log k) в среднем, но с памятью на весь массив. Замер — миллион случайных чисел, k от 10 до 100 000:

На Python heapq.nlargest находит десять наибольших из миллиона за 39 мс — в 10 раз быстрее, чем sorted; numpy.argpartition, который идёт по массиву numpy, а не по списку Python, ещё в шесть раз быстрее. При k = 100 000 куча уже втрое медленнее сортировки. На Go куча из десяти — 1,7 мс против 129 мс у slices.Sort; быстрый выбор держится около 7 мс почти при любом k и при k = 100 000 обгоняет кучу вчетверо.

Когда данные на нескольких серверах

Пусть данные разложены по нескольким машинам. Десять сильнейших землетрясений найти легко: каждая машина присылает свои десять сильнейших, и из присланного выбираем десять. Это точный ответ: землетрясение, которое входит в общую десятку, тем более входит в десятку своей машины. Совсем другое дело — самые частые слова, если каждая машина хранит свою часть текста. Частота слова — сумма по машинам, и слово, которое везде на одиннадцатом месте, может оказаться в целом на первом, а ни одна машина его не пришлёт.

Ниже «Война и мир», разложенная по пяти «серверам»: четыре тома и эпилог. Каждый том считает свои слова и присылает свои первые s вместе со счётом, а координатор складывает присланное.

При k = 20 и s = 20 потерялось слово «так». В романе оно встречается 1990 раз и стоит в общем топе, но в двадцатку своего тома попало только во втором томе, а в остальных четырёх не дотянуло до порога совсем немного: в эпилоге, например, 178 раз при пороге 181. Координатор насчитал ему 580. Его место заняло «из» с 887 присланными упоминаниями из 1848 настоящих. Если тома присылают по 40 слов, ответ становится точным. Но при k = 100 даже 160 слов от тома мало: теряются «ежели» и «граф», у которых в эпилоге и четвёртом томе мало упоминаний.

Что можно гарантировать? Том, который не прислал слово, встречает его не чаще, чем последнее присланное слово, — это его порог. Поэтому настоящий счёт слова не меньше присланной суммы и не больше её плюс сумма порогов тех томов, которые слово не прислали. Если нижняя граница слова выше верхних границ всех остальных кандидатов, оно в топе наверняка. Так делает Elasticsearch: агрегация terms по умолчанию просит у каждого шарда не k, а 1,5k + 10 значений и сообщает в ответе doc_count_error_upper_bound — оценку того, сколько могло потеряться. Точный ответ без второго прохода по данным в общем случае не получить, и это вторая задача в конце статьи.

Если же данных столько, что не помещается даже счётчик для каждого слова, точный топ уже не найти и на одной машине. Тогда частоты оценивают приближённо — наброском Count-Min с кучей или алгоритмом Мисры — Гриса, — а куча из k по-прежнему отбирает лидеров.

Где работает

Базы данных. Запрос ORDER BY … LIMIT N не сортирует всю таблицу. В PostgreSQL это появилось в 2007 году: в сообщении к изменению Том Лейн описывает кучу из N лучших строк, через которую просеиваются входные, — около M log N сравнений вместо M log M — и называет главным выигрышем не сравнения, а память: при малом N сортировке не приходится сбрасывать данные на диск. В плане запроса EXPLAIN ANALYZE такая сортировка подписана «top-N heapsort».

Библиотеки. heapq.nlargest и heapq.nsmallest в Python устроены как куча из k: при k = 1 они просто вызывают max или min, а если k не меньше длины данных — сортируют, и в документации сказано, что результат тот же, что у sorted(…)[:k]. В C++ std::partial_sort в libstdc++ строит кучу из первых k элементов, пропускает через неё остальные и сортирует саму кучу. В numpy argpartition делает то же быстрым выбором.

Поиск и рекомендации. Поисковик находит тысячи подходящих документов, а показать нужно десять самых релевантных; рекомендательная система ищет k ближайших соседей в пространстве признаков. Когда кандидатов миллионы, а k — десятки, куча из k — естественный способ собрать ответ.

Второй приз

В 1883 году Льюис Кэрролл — математик Чарльз Доджсон — выпустил брошюру «Lawn Tennis Tournaments: The True Method of Assigning Prizes with a Proof of the Fallacy of the Present Method». Его возмутил турнир на выбывание: победитель в нём действительно сильнейший, но второй приз достаётся финалисту, а второй по силе игрок мог проиграть победителю ещё в первом круге. Кэрролл предложил турнир, где игрок выбывает, только когда трое доказанно сильнее его.

Задача Кэрролла — это top-k при k = 2, и её решение известно: второй по силе обязательно проиграл победителю, потому что больше ему проиграть было некому. Значит, после турнира на выбывание, который стоит n − 1 матчей, достаточно устроить второй турнир среди проигравших победителю — их не больше ⌈log₂ n⌉. Всего n + ⌈log₂ n⌉ − 2 сравнений, и в 1964 году Кислицын доказал, что меньше нельзя. Дональд Кнут разбирает и брошюру Кэрролла, и эти оценки в третьем томе «Искусства программирования», в разделе о выборе с наименьшим числом сравнений. Куча из k на потоке до этого оптимума не дотягивает, но она проста, не требует всех данных сразу, и в жизни этого обычно достаточно.

Задачи

В первой задаче — самые частые слова «Войны и мира» с правилом для равных, на котором ошибается почти каждая первая куча. Во второй — топы по томам, как в третьей сцене: какие слова наверняка в общем топе, а какие — только может быть.