Куча и приоритетная очередь: как найти k-й максимум за O(n log k)

Зачем вам куча?

Представьте: на собеседовании вас просят найти k-й по величине элемент в неотсортированном массиве. Первая мысль — отсортировать и взять нужный элемент. Собеседующий хмурится: «А можно быстрее?» Оказывается, можно — с помощью структуры данных, которая называется куча (heap). Она позволяет за O(log k) получать минимальный или максимальный элемент, а главное — строить её можно за линейное время. В этой статье вы научитесь применять кучу для поиска top-k элементов и поймёте, почему это эффективнее сортировки.

Что такое куча?

Куча — это бинарное дерево, в котором каждый родитель больше (max-heap) или меньше (min-heap) своих детей. Хранится она в обычном массиве: корень — элемент с индексом 0, левый ребёнок — 2*i+1, правый — 2*i+2. Главное свойство: корень всегда минимальный (в min-heap) или максимальный (в max-heap), доступ к нему за O(1).

Зачем это нужно? Допустим, мы хотим найти k наибольших чисел. Если поддерживать min-heap размера k, то корень будет наименьшим среди текущих кандидатов. Проходим по массиву:

• Если в куче меньше k элементов — добавляем.
• Если куча полна, сравниваем новый элемент с корнем: если он больше — удаляем корень и добавляем новый.
• Иначе — игнорируем.

В итоге в куче останутся k наибольших чисел, а корень — k-й по величине. Каждая вставка/удаление стоит O(log k), всего n операций → O(n log k). Но главная магия — heapify: превращение массива в кучу за O(n). Это линейное время позволяет начинать с готовой кучи без сортировки.

Код: наивный подход vs куча

Наивная сортировка

def kth_largest_sort(nums, k):
    nums.sort()  # O(n log n)
    return nums[-k]  # k-й по величине

Работает, но сортирует все элементы, даже те, что не нужны.

Min-heap размера k

import heapq

def kth_largest_heap(nums, k):
    # Шаг 1: строим min-heap из первых k элементов
    min_heap = nums[:k]
    heapq.heapify(min_heap)  # O(k) — линейная сборка

    # Шаг 2: обрабатываем остальные
    for num in nums[k:]:
        if num > min_heap[0]:  # только если больше текущего k-го
            heapq.heapreplace(min_heap, num)  # удаляем корень и вставляем новое

    return min_heap[0]  # корень — k-й по величине

Почему это работает: инвариант кучи гарантирует, что корень — минимальный среди k хранимых элементов. Любой элемент, не больший корня, не может войти в top-k, поэтому мы его пропускаем. После прохода в куче — ровно k наибольших чисел.

Распространённая ошибка: забыть сделать heapify для начального среза. Если добавлять элементы через heappush по одному, получится O(k log k) — тоже неплохо, но теряется элегантная линейная гарантия.

Ещё один классический пример: слияние k отсортированных списков

def merge_k_lists(lists):
    min_heap = []
    # Инициализируем кучу первыми элементами каждого списка
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(min_heap, (lst[0], i, 0))

    result = []
    while min_heap:
        val, list_idx, elem_idx = heapq.heappop(min_heap)
        result.append(val)
        # Добавляем следующий элемент из того же списка
        if elem_idx + 1 < len(lists[list_idx]):
            nxt = lists[list_idx][elem_idx + 1]
            heapq.heappush(min_heap, (nxt, list_idx, elem_idx + 1))
    return result

Каждый pop/push — O(log k), всего N элементов → O(N log k). Куча всегда даёт наименьший текущий головной элемент, гарантируя глобально отсортированный результат.

Ошибка: добавление следующего элемента до извлечения текущего может нарушить порядок. Всегда сначала pop, потом push.

Где ещё пригодится куча

• Поиск top-k в потоке данных — без хранения всего потока.
• Алгоритм Дейкстры — приоритетная очередь извлекает следующую вершину за O(log V).
• Планировщик задач — симуляция очередей с приоритетами.
• Скользящая медиана — две кучи (max-heap для нижней половины, min-heap для верхней).

Одна и та же идея — держать экстремальный элемент наверху и обновлять за логарифмическое время — решает десятки задач. Освоив кучу, вы начнёте замечать, где можно заменить O(n log n) на O(n log k) или даже O(n).

Что делать прямо сейчас

Возьмите случайный список чисел и число k. Вычислите k-й наибольший элемент без полной сортировки. Сначала реализуйте наивный вариант, затем перепишите с min-heap. Замерьте время на массиве из миллиона элементов — разница вас удивит. Если хотите усложнить, адаптируйте решение для потока чисел в реальном времени (например, для live-таблицы лидеров).

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