Куча в Python: алгоритм поиска k наибольших элементов

Зачем вам нужна куча в Python?

Вы готовитесь к собеседованию на позицию backend-разработчика, и интервьюер задаёт классическую задачу: «Дан неотсортированный массив из n целых чисел, верните k наибольших элементов». Первая мысль — отсортировать весь массив и взять срез: O(n log n) кажется приемлемым, но интервьюер явно ждёт более эффективного решения. Вы чувствуете, что можно избежать полной сортировки, когда нужны только top k. Это как смотреть на запертую дверь, имея связку ключей, среди которых есть нужный, но вы не знаете, какой именно. Знакомо? Тогда давайте разберёмся, как куча (heap) помогает решить эту задачу за O(n log k) и даже быстрее.

Что такое куча и как её построить за O(n)

Куча — это двоичное дерево, хранящееся в массиве, где каждый родительский узел упорядочен относительно своих детей: в min-куче родитель ≤ детей, в max-куче родитель ≥ детей. Ключевой момент: если взять произвольный массив и «хипифицировать» его, начиная с последнего внутреннего узла и опуская элементы вниз, мы получим корректную кучу за O(n), а не за O(n log n).

Почему heapify работает за линейное время?

Нижние уровни дерева содержат много узлов, но их высота мала, поэтому стоимость опускания вниз невелика. Верхние уровни имеют мало узлов, но элементы могут спускаться далеко, однако их вклад ограничен. Суммируя затраты по всем уровням, получаем геометрическую прогрессию, которая сходится к O(n).

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

• Построить min-кучу размера k из первых k элементов (O(k)).
• Для каждого оставшегося элемента: если он больше корня кучи, заменяем корень и опускаем его вниз (O(log k)).

Итого сложность: O(k + (n-k) log k) ≈ O(n log k) в худшем случае, но при малых k относительно n это почти линейно. Главное — начальное построение кучи для всего массива — O(n), что является основой для многих других алгоритмов на куче (например, пирамидальная сортировка или поиск медианы).

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

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

def k_largest_sort(nums, k):
    # Время O(n log n), память O(n) для Timsort в CPython
    return sorted(nums, reverse=True)[:k]

Просто, но расточительно, когда n огромно, а k мало.

Эффективный способ: куча

import heapq

def k_largest_heap(nums, k):
    """
    Возвращает k наибольших элементов, используя min-кучу размера k.
    Время: O(n log k) (операции push/pop) + O(k) на построение начальной кучи
    Память: O(k)
    """
    if k == 0:
        return []
    # Шаг 1: строим min-кучу из первых 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)  # заменяем корень и восстанавливаем кучу за O(log k)
    # В куче теперь k наибольших элементов, но они не отсортированы
    return sorted(min_heap, reverse=True)  # опционально: возвращаем по убыванию

Почему это работает: куча всегда содержит k самых больших элементов, встреченных на данный момент. Любой элемент, меньший, чем минимальный в куче, не может повлиять на ответ, поэтому мы его пропускаем. heapreplace эффективнее отдельного pop и push, так как избегает лишних операций восстановления.

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

• Забыть heapify для начального среза. Если просто присвоить min_heap = nums[:k] и начать вставлять/удалять без heapify, список не будет кучей, и операции станут O(k), вырождаясь в O(nk).
• Использовать max-кучу, когда нужна min-куча. В Python heapq реализует только min-кучу. Для эмуляции max-кучи инвертируйте знак (-num). Путаница с инверсией приводит к ошибкам на границах.

Проверим на примере из собеседования:

>>> nums = [3, 1, 5, 12, 2, 11, 7]
>>> k_largest_heap(nums, 3)
[12, 11, 7]

Второй вариант задачи: k-й наибольший элемент

Часто спрашивают: «Верните k-й наибольший элемент». Используя ту же кучу, мы можем остановиться после обработки всех элементов и просто взять корень кучи:

def kth_largest(nums, k):
    min_heap = nums[:k]
    heapq.heapify(min_heap)
    for num in nums[k:]:
        if num > min_heap[0]:
            heapq.heapreplace(min_heap, num)
    return min_heap[0]  # k-й наибольший

Сложность та же, но мы избегаем финальной сортировки — идеально, когда нужно только одно значение.

Почему это знание меняет подход к задачам

Кучи превращают задачу «нужны top k» из задачи сортировки в задачу динамического отбора. Представьте, что вы обрабатываете поток данных с датчиков и постоянно хотите знать 10 самых высоких температур. Min-куча размера 10 даёт O(log 10) ≈ O(1) на каждое показание, с постоянной памятью — сортировка такого не может.

Помимо собеседований, кучи используются в:

• Алгоритме Дейкстры (приоритетная очередь для извлечения минимума)
• Симуляции событий (следующая временная метка)
• K-way merge (эффективное слияние k отсортированных списков)

Понимание того, что построение кучи линейно, а не логарифмическое, позволяет замечать возможности сократить алгоритм с O(n log n) до O(n + k log n) или даже O(n), когда k — константа. Это инструмент, который превращает «грубую силу» в «элегантность» всего несколькими строками кода.

Ваш ход

Выберите задачу, которую недавно решали и в которой использовали сортировку или поиск экстремумов. Попробуйте переформулировать её через кучу — возможно, удастся сократить логарифмический множитель или уменьшить использование памяти. Поделитесь решением или вопросом в комментариях — интересно увидеть, как вы примените это новое знание в своих задачах!

Удачи в кодинге! 🚀