Поиск K-го наименьшего элемента через кучу за O(n log k)

Зачем искать K-й наименьший элемент быстрее, чем сортировкой?

Представьте: на собеседовании вам задают найти K-й наименьший элемент в неотсортированном массиве. Первая мысль — отсортировать всё и взять индекс K-1. Это работает, но за O(n log n), хотя нужен всего один элемент. Есть ли способ умнее? Да, с помощью кучи мы можем решить задачу за O(n log k), что особенно эффективно, когда K много меньше N. В этой статье вы освоите этот алгоритм на Python и узнаете, как он применяется в реальных задачах и на платформах вроде LeetCode.

Идея алгоритма: храним только K наименьших

Ключевое прозрение: нам не нужно поддерживать весь массив отсортированным, достаточно знать K наименьших значений. Если мы будем хранить структуру, которая всегда даёт наибольший из этих K элементов, то всё, что больше, можно отбрасывать. Для этого идеально подходит max-куча (в Python — через отрицательные значения, потому что heapq реализует min-кучу).

Почему это работает?

• Куча хранит не более K элементов.
• Её корень — наибольший среди K наименьших, которые мы видели.
• Для нового числа x: если куча заполнена, сравниваем x с корнем. Если x меньше корня, то корень не может быть в финальном множестве K наименьших — заменяем его. Если x больше или равен, игнорируем.
• После обработки всех элементов корень кучи — это K-й наименьший элемент.

Построение кучи из массива занимает O(n), а каждая вставка/удаление — O(log k). Суммарно получаем O(n log k), что при малом K значительно быстрее сортировки.

Реализация на Python: от простого к эффективному

Наивный способ: сортировка

def kth_smallest_sort(nums, k):
    return sorted(nums)[k-1]  # O(n log n) по времени, O(n) по памяти

Просто, но для больших массивов логарифмический фактор вредит.

Решение на куче

import heapq

def kth_smallest_heap(nums, k):
    # heapq в Python — min-куча, поэтому храним отрицательные значения для имитации max-кучи
    max_heap = []
    for num in nums:
        if len(max_heap) < k:
            heapq.heappush(max_heap, -num)  # отрицательное значение для max-кучи
        else:
            # если текущее число меньше наибольшего в куче, заменяем его
            if -num > max_heap[0]:  # max_heap[0] — наименьшее отрицательное = наибольшее исходное
                heapq.heapreplace(max_heap, -num)
    return -max_heap[0]  # корень max-кучи (как отрицательное) — K-й наименьший

Почему это магия: мы не сортируем весь массив, а держим лишь окно размера K. Операция heapreplace заменяет корень и добавляет новый элемент за один O(log k), избегая лишних операций.

Типичные ошибки и как их избежать

• Забыть про отрицательные значения — куча будет вести себя как min-куча, и вы получите K-й наибольший вместо наименьшего. Всегда храните -num для max-кучи.
• Использовать heappush и heappop отдельно — это две операции O(log k), когда можно обойтись одной heapreplace. Используйте heapreplace, когда куча уже заполнена.
• Неправильное направление сравнения — можно случайно оставить большие элементы и потерять настоящий K-й наименьший. Помните: с отрицательными числами большее исходное число становится меньшим отрицательным.

Вариация: K-й наибольший элемент

Тот же приём работает для поиска K-го наибольшего: используем min-кучу размера K, хранящую K наибольших элементов. Корень будет ответом. Код почти идентичен, только без отрицаний.

def kth_largest_heap(nums, k):
    min_heap = []
    for num in nums:
        if len(min_heap) < k:
            heapq.heappush(min_heap, num)
        else:
            if num > min_heap[0]:
                heapq.heapreplace(min_heap, num)
    return min_heap[0]

Обе задачи часто встречаются на LeetCode (например, 215. Kth Largest Element in an Array) и отлично демонстрируют понимание куч.

Почему это важно для вашего роста

Освоив этот подход, вы меняете мышление: вместо полной сортировки вы спрашиваете, какую минимальную информацию нужно сохранить. Это умение применяется в потоковых данных (поиск медианы), планировании задач по приоритетам и даже в алгоритме Дейкстры, где приоритетная очередь — это куча.

Выигрыш в производительности реален: для N = 1 миллион и K = 10 решение с кучей делает примерно 10 × log 10 ≈ 33 операций на элемент против 20 для полной сортировки (log N ≈ 20), но с гораздо меньшим перемещением данных и лучшей кэш-эффективностью. На собеседовании демонстрация перехода от O(n log n) к O(n log k) (или O(n) на heapify + O(n log k)) показывает, что вы понимаете компромиссы, а не просто заучиваете паттерны.

Ваше задание

Возьмите неотсортированный список, выберите K и попробуйте закодировать версию с max-кучей, не подглядывая в решение выше. Протестируйте её против наивной сортировки на случайных массивах и наблюдайте, как разница в скорости растёт при уменьшении K относительно N.

Челлендж: модифицируйте код для корректной обработки дубликатов и возврата K-го различного наименьшего элемента. Как изменится размер кучи?

Удачного кодинга, и пусть ваши кучи всегда остаются сбалансированными!