← Назад к курсу

Урок по работе с кучей (heap) в Python

Типовые задачи на собеседованиях разной сложности

1. Что такое куча (heap)?

Куча — это специальная древовидная структура данных, которая удовлетворяет свойству кучи:

  • Min-Heap: значение родителя меньше или равно значениям его потомков
  • Max-Heap: значение родителя больше или равно значениям его потомков

В Python для работы с кучами используется модуль heapq.

import heapq

# Создание кучи из списка
arr = [5, 7, 9, 1, 3]
heapq.heapify(arr)  # Преобразует список в кучу за O(n)
print(f"Куча: {list(arr)}")  # [1, 3, 9, 7, 5]

# Добавление элемента
heapq.heappush(arr, 4)
print(f"После добавления 4: {list(arr)}")

# Извлечение минимального элемента
min_val = heapq.heappop(arr)
print(f"Минимальный элемент: {min_val}")
print(f"Куча после извлечения: {list(arr)}")

2. Основные операции с кучей

import heapq

class HeapDemo:
    def __init__(self):
        self.heap = []
    
    def demo_operations(self):
        print("=== Основные операции с кучей ===")
        
        # Добавление элементов
        for num in [10, 5, 20, 3, 8]:
            heapq.heappush(self.heap, num)
            print(f"Добавили {num}: {self.heap}")
        
        # Просмотр минимального элемента без удаления
        print(f"\nМинимальный элемент (не удаляя): {self.heap[0]}")
        
        # Извлечение элементов в отсортированном порядке
        print("\nИзвлечение элементов в порядке возрастания:")
        while self.heap:
            print(f"  Извлекли: {heapq.heappop(self.heap)}, осталось: {self.heap}")
        
        # Max-heap через отрицательные значения
        print("\n=== Max-Heap через отрицательные значения ===")
        max_heap = []
        for num in [10, 5, 20, 3, 8]:
            heapq.heappush(max_heap, -num)
        print(f"Max-heap (с отрицательными): {max_heap}")
        print(f"Максимальный элемент: {-max_heap[0]}")

demo = HeapDemo()
demo.demo_operations()

3. Задачи начального уровня

Задача 1: K-й наименьший элемент

def kth_smallest(arr, k):
    """
    Найти k-й наименьший элемент в массиве.
    Сложность: O(n + k*logn)
    """
    heapq.heapify(arr)  # O(n)
    
    result = None
    for _ in range(k):  # O(k*logn)
        result = heapq.heappop(arr)
    
    return result

# Альтернатива с nsmallest
def kth_smallest_builtin(arr, k):
    """Использование встроенной функции heapq.nsmallest"""
    return heapq.nsmallest(k, arr)[-1]

# Пример использования
arr = [7, 10, 4, 3, 20, 15]
k = 3
print(f"Массив: {arr}")
print(f"{k}-й наименьший элемент (heap): {kth_smallest(arr.copy(), k)}")
print(f"{k}-й наименьший элемент (nsmallest): {kth_smallest_builtin(arr, k)}")

Задача 2: Объединение K отсортированных списков

def merge_k_sorted_lists(lists):
    """
    Объединить k отсортированных списков в один отсортированный список.
    Сложность: O(N*logk), где N - общее количество элементов
    """
    heap = []
    result = []
    
    # Инициализируем кучу первыми элементами каждого списка
    for i, lst in enumerate(lists):
        if lst:  # если список не пустой
            heapq.heappush(heap, (lst[0], i, 0))  # (значение, индекс списка, индекс в списке)
    
    while heap:
        val, list_idx, element_idx = heapq.heappop(heap)
        result.append(val)
        
        # Добавляем следующий элемент из того же списка
        if element_idx + 1 < len(lists[list_idx]):
            next_val = lists[list_idx][element_idx + 1]
            heapq.heappush(heap, (next_val, list_idx, element_idx + 1))
    
    return result

# Пример
lists = [[1, 4, 5], [1, 3, 4], [2, 6]]
merged = merge_k_sorted_lists(lists)
print(f"\nОбъединение списков {lists}:")
print(f"Результат: {merged}")

4. Задачи среднего уровня

Задача 3: Топ K частых элементов

from collections import Counter
import heapq

def top_k_frequent(nums, k):
    """
    Найти k самых частых элементов.
    Сложность: O(n*logk)
    """
    # Считаем частоту каждого элемента
    freq = Counter(nums)
    
    # Создаем min-heap размером k
    heap = []
    
    for num, count in freq.items():
        if len(heap) < k:
            heapq.heappush(heap, (count, num))
        elif count > heap[0][0]:
            heapq.heapreplace(heap, (count, num))
    
    # Извлекаем элементы из кучи
    return [num for _, num in heap]

# Пример
nums = [1, 1, 1, 2, 2, 3, 4, 4, 4, 4]
k = 2
print(f"\nТоп-{k} самых частых элементов в {nums}:")
print(f"Результат: {top_k_frequent(nums, k)}")

Задача 4: Поиск медианы в потоке данных

class MedianFinder:
    """
    Класс для нахождения медианы в потоке данных.
    Использует две кучи: max-heap для левой половины и min-heap для правой.
    """
    def __init__(self):
        self.max_heap = []  # левая половина (max-heap через отрицательные значения)
        self.min_heap = []  # правая половина (min-heap)
    
    def add_num(self, num):
        # Добавляем в левую кучу
        heapq.heappush(self.max_heap, -num)
        
        # Балансировка: максимальный элемент из левой кучи
        # должен быть <= минимального из правой
        if (self.max_heap and self.min_heap and 
            (-self.max_heap[0]) > self.min_heap[0]):
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        
        # Балансировка размеров
        if len(self.max_heap) > len(self.min_heap) + 1:
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        elif len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))
    
    def find_median(self):
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        return (-self.max_heap[0] + self.min_heap[0]) / 2

# Пример
finder = MedianFinder()
numbers = [5, 15, 1, 3, 10, 20, 2]
print("\nПоиск медианы в потоке данных:")
for num in numbers:
    finder.add_num(num)
    print(f"Добавили {num}, медиана: {finder.find_median():.2f}")

5. Продвинутые задачи

Задача 5: Task Scheduler

from collections import Counter
import heapq

def least_interval(tasks, n):
    """
    Минимальное время выполнения задач с одинаковыми задачами,
    разделенными интервалом охлаждения.
    
    Args:
        tasks: список задач (символы)
        n: интервал охлаждения между одинаковыми задачами
    
    Сложность: O(N*logk), где k - количество уникальных задач
    """
    # Считаем частоту задач
    freq = Counter(tasks)
    
    # Создаем max-heap из частот (через отрицательные значения)
    max_heap = [-count for count in freq.values()]
    heapq.heapify(max_heap)
    
    time = 0
    queue = []  # очередь для задач на охлаждении
    
    while max_heap or queue:
        time += 1
        
        if max_heap:
            # Берем задачу с максимальной частотой
            count = heapq.heappop(max_heap) + 1  # +1 потому что отрицательное
            
            if count < 0:  # если остались еще такие задачи
                # Добавляем в очередь с временем, когда будет доступна
                queue.append((count, time + n))
        
        # Проверяем, не истекло ли время охлаждения
        if queue and queue[0][1] == time:
            heapq.heappush(max_heap, queue.pop(0)[0])
    
    return time

# Пример
tasks = ["A", "A", "A", "B", "B", "C"]
n = 2
result = least_interval(tasks, n)
print(f"\nTask Scheduler: задачи {tasks}, интервал {n}")
print(f"Минимальное время выполнения: {result}")
print("Одно из возможных расписаний: A -> B -> C -> A -> B -> A")

Задача 6: K отсортированных массивов - найти диапазон

def smallest_range(nums):
    """
    Найти наименьший диапазон, который включает хотя бы по одному
    элементу из каждого из k отсортированных списков.
    
    Сложность: O(N*logk), где N - общее количество элементов
    """
    heap = []
    current_max = float('-inf')
    
    # Инициализируем кучу первыми элементами каждого списка
    for i, arr in enumerate(nums):
        heapq.heappush(heap, (arr[0], i, 0))
        current_max = max(current_max, arr[0])
    
    start, end = float('-inf'), float('inf')
    
    while heap:
        current_min, list_idx, element_idx = heapq.heappop(heap)
        
        # Обновляем диапазон
        if current_max - current_min < end - start:
            start, end = current_min, current_max
        
        # Если в текущем списке есть еще элементы
        if element_idx + 1 < len(nums[list_idx]):
            next_val = nums[list_idx][element_idx + 1]
            heapq.heappush(heap, (next_val, list_idx, element_idx + 1))
            current_max = max(current_max, next_val)
        else:
            # Один из списков закончился
            break
    
    return [start, end]

# Пример
k_sorted_arrays = [
    [4, 10, 15, 24, 26],
    [0, 9, 12, 20],
    [5, 18, 22, 30]
]

range_result = smallest_range(k_sorted_arrays)
print(f"\nK отсортированных массивов: {k_sorted_arrays}")
print(f"Наименьший диапазон, покрывающий все массивы: {range_result}")

6. Практические упражнения

def practice_exercises():
    """Практические упражнения для закрепления материала"""
    
    print("=== Практические упражнения ===")
    
    # Упражнение 1: Реализовать Max-Heap
    class MaxHeap:
        def __init__(self):
            self.heap = []
        
        def push(self, val):
            heapq.heappush(self.heap, -val)
        
        def pop(self):
            return -heapq.heappop(self.heap)
        
        def peek(self):
            return -self.heap[0] if self.heap else None
    
    max_heap = MaxHeap()
    for num in [3, 1, 6, 5, 2, 4]:
        max_heap.push(num)
    
    print("Max-Heap элементы в порядке убывания:")
    while max_heap.peek() is not None:
        print(f"  {max_heap.pop()}")
    
    # Упражнение 2: Сортировка почти отсортированного массива
    def sort_k_sorted(arr, k):
        """
        Сортировка массива, в котором каждый элемент
        находится не более чем в k позициях от своей правильной позиции.
        """
        heap = arr[:k+1]
        heapq.heapify(heap)
        
        result = []
        for i in range(k+1, len(arr)):
            result.append(heapq.heappop(heap))
            heapq.heappush(heap, arr[i])
        
        while heap:
            result.append(heapq.heappop(heap))
        
        return result
    
    arr = [6, 5, 3, 2, 8, 10, 9]
    k = 3
    print(f"\nСортировка почти отсортированного массива (k={k}):")
    print(f"Исходный: {arr}")
    print(f"Отсортированный: {sort_k_sorted(arr, k)}")

practice_exercises()

7. Полезные советы для собеседований

  1. Когда использовать кучу:

    • Нужно часто получать min/max элемент
    • Работа с K наибольшими/наименьшими элементами
    • Объединение отсортированных списков
    • Реализация очереди с приоритетом
  2. Временная сложность:

    • heapify: O(n)
    • heappush: O(logn)
    • heappop: O(logn)
    • heap[0]: O(1) для получения min элемента
  3. Max-Heap в Python:

    • Нет встроенной поддержки max-heap
    • Используйте отрицательные значения: heapq.heappush(heap, -value)
    • Или создайте свой класс с __lt__ методом
  4. Частые ошибки:

    • Не используйте heap как обычный список (индексы нарушают свойства кучи)
    • После модификации элементов нужно вызывать heapify заново
    • Не путайте heapreplace (pop → push) и heappushpop (push → pop)

8. Дополнительные ресурсы

def additional_resources():
    print("\n=== Дополнительные ресурсы ===")
    resources = [
        "1. Официальная документация heapq: https://docs.python.org/3/library/heapq.html",
        "2. Визуализация работы кучи: https://visualgo.net/en/heap",
        "3. LeetCode теги для практики: Heap, Priority Queue",
        "4. Книга: 'Алгоритмы. Построение и анализ' - Томас Кормен (глава 6)",
        "5. Курс: Algorithms, Part I на Coursera (Princeton)"
    ]
    for resource in resources:
        print(resource)

additional_resources()

Заключение

Кучи — мощная структура данных для задач, связанных с приоритетами и выбором экстремальных значений. На собеседованиях задачи на кучу проверяют понимание:

  1. Основных операций и их сложности
  2. Применения в реальных задачах (системы очередей, планировщики)
  3. Комбинации с другими структурами данных
  4. Умение оптимизировать решения

Ключевые моменты для запоминания:

  • В Python используется min-heap
  • Max-heap реализуется через отрицательные значения
  • heapq модифицирует список на месте
  • Для частых задач есть встроенные функции (nlargest, nsmallest)

Практикуйтесь на задачах с LeetCode (теги: Heap, Priority Queue) для лучшего понимания!