Примеры использования кучи в Python: от сортировки до алгоритмов графов

Введение
Куча (heap) — это эффективная структура данных, позволяющая быстро получать доступ к элементу с максимальным или минимальным значением. В Python для работы с кучей используется модуль heapq, реализующий мин-кучу. В этой статье мы рассмотрим ключевые сценарии применения кучи в реальных задачах, сопровождая их примерами кода.

---

1. Основы работы с кучей в Python

Куча — это двоичное дерево, где каждый родительский узел меньше (мин-куча) или больше (макс-куча) своих дочерних узлов. В heapq реализована мин-куча, поэтому корень всегда содержит наименьший элемент.

Основные операции:

• heappush(heap, element): Добавление элемента.
• heappop(heap): Извлечение наименьшего элемента.
• heapify(list): Преобразование списка в кучу.
• nlargest(k, iterable), nsmallest(k, iterable): Поиск k наибольших/наименьших элементов.

---

2. Примеры использования

2.1 Нахождение N наибольших/наименьших элементов

Модуль heapq предоставляет функции nsmallest и nlargest, которые эффективнее сортировки при работе с большими данными.

import heapq

data = [3, 1, 4, 1, 5, 9, 2, 6]
print("3 наименьших:", heapq.nsmallest(3, data))  # [1, 1, 2]
print("3 наибольших:", heapq.nlargest(3, data))   # [9, 6, 5]

2.2 Реализация приоритетной очереди

Приоритетная очередь обрабатывает элементы в порядке их приоритета. Пример реализации:

import heapq

class PriorityQueue:
    def __init__(self):
        self._heap = []
        self._index = 0  # Для устранения коллизий при равных приоритетах

    def push(self, item, priority):
        heapq.heappush(self._heap, (priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._heap)[-1]

# Использование
pq = PriorityQueue()
pq.push("задача 1", 3)
pq.push("задача 2", 1)
print(pq.pop())  # "задача 2" (высший приоритет)

2.3 Алгоритм Дейкстры для поиска кратчайшего пути

Куча используется для эффективного выбора вершины с минимальным расстоянием.

import heapq

def dijkstra(graph, start):
    distances = {node: float('inf') for node in graph}
    distances[start] = 0
    heap = [(0, start)]
    
    while heap:
        current_dist, current_node = heapq.heappop(heap)
        if current_dist > distances[current_node]:
            continue
        for neighbor, weight in graph[current_node].items():
            distance = current_dist + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(heap, (distance, neighbor))
    return distances

# Пример графа
graph = {
    'A': {'B': 2, 'C': 5},
    'B': {'C': 1, 'D': 3},
    'C': {'D': 4},
    'D': {}
}
print(dijkstra(graph, 'A'))  # {'A': 0, 'B': 2, 'C': 3, 'D': 5}

2.4 Сортировка с помощью кучи (Heapsort)

Пример реализации пирамидальной сортировки:

import heapq

def heapsort(arr):
    heapq.heapify(arr)
    return [heapq.heappop(arr) for _ in range(len(arr))]

data = [3, 1, 4, 1, 5, 9, 2, 6]
sorted_data = heapsort(data)
print(sorted_data)  # [1, 1, 2, 3, 4, 5, 6, 9]

2.5 Слияние отсортированных последовательностей

Куча позволяет объединить k отсортированных списков за O(n log k):

import heapq

def merge_sorted_arrays(arrays):
    heap = []
    for i, arr in enumerate(arrays):
        if arr:
            heapq.heappush(heap, (arr[0], i, 0))
    
    result = []
    while heap:
        val, arr_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        if elem_idx + 1 < len(arrays[arr_idx]):
            next_val = arrays[arr_idx][elem_idx + 1]
            heapq.heappush(heap, (next_val, arr_idx, elem_idx + 1))
    return result

arrays = [[1, 3, 5], [2, 4, 6], [0, 7]]
print(merge_sorted_arrays(arrays))  # [0, 1, 2, 3, 4, 5, 6, 7]

2.6 Управление событиями в симуляциях

Куча помогает обрабатывать события в порядке их времени:

import heapq

class EventSimulator:
    def __init__(self):
        self.events = []
    
    def schedule_event(self, time, event):
        heapq.heappush(self.events, (time, event))
    
    def run(self):
        while self.events:
            time, event = heapq.heappop(self.events)
            print(f"Время {time}: обработано событие {event}")

sim = EventSimulator()
sim.schedule_event(5, "завершение задачи")
sim.schedule_event(2, "проверка состояния")
sim.run()
# Вывод:
# Время 2: обработано событие проверка состояния
# Время 5: обработано событие завершение задачи

---

3. Особенности и ограничения

• Макс-куча: Для реализации используйте инвертированные значения: heapq.heappush(heap, -x).
• Обновление приоритета: heapq не поддерживает напрямую. Решение — добавлять новые элементы и игнорировать устаревшие.
• Производительность: Операции heappush и heappop имеют сложность O(log n), heapify — O(n).

---

Заключение

Кучи в Python находят применение в задачах, требующих эффективного доступа к экстремальным элементам: сортировка, приоритетные очереди, алгоритмы на графах, управление событиями. Модуль heapq предоставляет инструменты для работы с кучей, но требует внимания при обновлении элементов и реализации макс-кучи. Использование кучи оптимизирует выполнение операций и снижает сложность алгоритмов, делая их применимыми для больших данных.