heapq и max heap в Python

Язык программирования Python предоставляет множество инструментов для хранения и представления данных: int, float, bool, complex, list, dict, tuple, set, frozenset и множество других.

И если с базовыми типами все достаточно просто, то более абстрактные сущности зачастую вызывают вопросы — они требуют больше теоретических данных как для понимания, так и для практического применения.

Одна из таких структур предлагает нестандартный способ компоновки данных и еще более нестандартные способы работы с ними. Такая структура называется кучей. Именно о ней и пойдет речь в этой статье.

Обратите внимание, что все демонстрируемые примеры запускались с помощью интерпретатора Python версии 3.10.12, установленном на облачном сервере Timeweb Cloud под управлением операционной системы Ubuntu 22.04.

Каждый скрипт размещался в отдельном файле с расширением .py (например, script.py), после чего запускался с помощью команды интерпретатора Python:

python script.py

cloud

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

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

Проще говоря, куча — это бинарное дерево, узлы которого представлены в виде линейной последовательности.

Схема последовательного расположения узлов двоичного дерева кучи внутри массива (источник: Блог программиста).

Концепт кучи в Питоне (да и вообще в программировании) связан с рядом фундаментальных понятий:

• Бинарное дерево поиска (binary search tree, BST). Структура данных, каждый родительский узел которого (начиная с корневого) рекурсивно имеет только два дочерних узла.
• Свойство кучи (heap property). Правило, которому должен удовлетворять каждый узел двоичной кучи. Бывает двух видов:
  • Min-куча. Значение в любом родительском узле меньше или равно значениям его детей.
  • Max-куча. Значение в родительском узле больше или равно значениям его детей.
• Приоритетная очередь (priority queue). Структура данных в виде очереди, в которой каждый элемент имеет приоритет (важность) — извлекается либо наиболее важный, либо наименее важный, а не тот, который был вставлен первым.
• Извлечение элемента (heappop). Получение либо минимального, либо максимального элемента из кучи — то есть корня.
• Вставка элемента (heappush). Добавление нового элемента либо в минимальную, либо в максимальную позицию кучи.
• Всплытие (sift-up). Процесс поднятия нового элемента до тех пор, пока его позиция не станет соответствовать свойству кучи.
• Утопление (sift-down). Процесс опускания элемента до тех пор, пока его позиция не станет соответствовать свойству кучи.
• Сортировка кучей (heap sort). Сортировка списка таким образом, чтобы конечные позиции его элементов соответствовали свойству кучи.

С теоретической точки зрения куча может показаться сложным и многокомпонентным концептом, учитывая количество явно обозначенных понятий, которые с ней связаны. Однако на практике базовое понятие двоичного дерева проявляется достаточно просто.

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

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

Приоритетная очередь

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

По сути, куча гарантирует, что следующей будет взята задача с наивысшим (или наименьшим) приоритетом.

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

Дискретно-событийное моделирование

В системах моделирования (сетевые симуляторы, игровые движки) каждое событие имеет метку времени, а обработка начинается с наиболее ранних.

Куча позволяет организовать упорядоченную «таймерную» очередь, состоящую из событий с соответствующими временными метками.

Такая система гарантирует детерминированный порядок и высокую производительность во время динамической вставки и удаления по ходу моделирования.

Менеджмент памяти и буферизация

В операционных системах и базах данных кучи применяются для управления кэшем или страницами памяти.

Также куча встречается в реализациях сетевых буферов и дисковых очередей, где важен приоритет передачи или время прихода пакета — здесь также нужно быстро извлекать крайние элементы.

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

Модуль heapq для работы с кучей

В языке программирования Python для работы с кучей есть специальный модуль, доступный из стандартной библиотеки — heapq.

По умолчанию куча в модуле heapq является минимальной (min-heap) — корневой элемент с индексом 0 всегда наименьший, а каждый родительский элемент всегда меньше или равен дочерним. Благодаря этому можно эффективно получать и удалять минимальный элемент.

Для получения максимальной (max-heap) кучи используют отрицательные значения её узлов (элементов). Таким образом, с помощью инвертирования из минимальной кучи можно получить максимальную.

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

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7] # создание списка с исходными данными для кучи
print(someNumbers) # ВЫВОД: [8, 3, 5, 1, 6, 2, 4, 7]

heapq.heapify(someNumbers) # превращение список в кучу
print(someNumbers) # ВЫВОД: [1, 3, 2, 7, 6, 5, 4, 8]

# бинарное дерево приобретает следующий вид:
               1
        3            2
    7     6     5     4
8

heapq.heappush(someNumbers, 0) # добавление нового элемента в кучу
print(someNumbers) # ВЫВОД: [0, 1, 2, 3, 6, 5, 4, 8, 7]

# бинарное дерево приобретает следующий вид:
               0
        1            2
    3     6     5     4
8     7

someValue = heapq.heappop(someNumbers) # извлечение минимального элемента из кучи
print(someNumbers) # ВЫВОД: [1, 3, 2, 7, 6, 5, 4, 8]

# бинарное дерево приобретает следующий вид:
               1
        3            2
    7     6     5     4
8

print(someValue) # ВЫВОД: 0

В показанном примере выполняются следующие действия:

• Элементы списка сортируются таким образом, чтобы соответствовать свойству кучи. Для этого используется метод heapify(). 
• В отсортированный список добавляется новый элемент на позицию, соответствующую свойству кучи. Для этого используется метод heappush().
• Из отсортированного списка извлекается и возвращается корень кучи (элемент с наименьшим значением), а последний элемент перемещается на позицию, соответствующую свойству кучи. Для этого используется метод heappop().

Как можно заметить, библиотека heapq оперирует над итерируемыми объектами встроенного типа list, являясь своего рода инструментом для модификации списков. Другие типы итерируемых объектов куча не поддерживает.

Методы heapq для работы с кучей

heapify(список)

Преобразует список в кучу за счет расположения элементов согласно свойству кучи (min-heap):

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7]
heapq.heapify(someNumbers)

print(someNumbers) # ВЫВОД: [1, 3, 2, 7, 6, 5, 4, 8]

Важно понимать, что куча не сортирует элементы идеально точно — она лишь следует свойству, при котором каждый дочерний элемент больше родительского:

import heapq

firstSorted = [1, 3, 2, 7, 6, 5, 4, 8] # заранее отсортированный список кучи
secondSorted = [1, 2, 3, 5, 6, 7, 4, 8] # тот же самый отсортированный список кучи, но с немного другим порядком элементов —  пара 3 и 2, равно как и пара 7 и 5, поменяны местами

heapq.heapify(secondSorted)

print(firstSorted) # ВЫВОД: [1, 3, 2, 7, 6, 5, 4, 8]
print(secondSorted) # ВЫВОД: [1, 2, 3, 5, 6, 7, 4, 8]

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

Таким образом, если представить два коротких списка кучи:

first = [1, 4, 9]
second = [1, 9, 4]

То оба эти списка будут удовлетворять свойству кучи — и первая, и вторая запись полностью корректны.

heappush(список, элемент)

Добавляет новый элемент в кучу, заставляя его всплать (sift-up) позицию, соответствующую свойству кучи (min-heap):

import heapq

someNumbers = [1, 3, 2, 7, 6, 5, 4, 8]
heapq.heappush(someNumbers, 0)
print(someNumbers) # ВЫВОД: [0, 1, 2, 3, 6, 5, 4, 8, 7]

heappop(список) : элемент

Извлекает и возвращает корень кучи (элемент с наименьшим значением), а последний элемент утапливается (sift-down) вниз — в позицию, соответствующую свойству кучи (min-heap):

import heapq

someNumbers = [1, 3, 2, 7, 6, 5, 4, 8]
heapq.heappush(someNumbers, 0)
print(someNumbers) # ВЫВОД: [0, 1, 2, 3, 6, 5, 4, 8, 7]

heappushpop(список, элемент) : элемент

Добавляет новый элемент и извлекает корень кучи:

import heapq

someNumbers = [1, 3, 2, 7, 6, 5, 4, 8]

heapq.heappushpop(someNumbers, 0) # вставляет 0 и извлекает 0
print(someNumbers) # ВЫВОД: [1, 3, 2, 7, 6, 5, 4, 8]

heapq.heappushpop(someNumbers, 20) # вставляет 20 и извлекает 1
print(someNumbers) # ВЫВОД: [2, 3, 4, 7, 6, 5, 20, 8]

Таким образом, если во время вставки новый элемент становится корнем кучи, то он сразу же извлекается обратно.

heapreplace(список, элемент) : элемент

Извлекает корень кучи и добавляет новый элемент:

import heapq

someNumbers = [1, 3, 2, 7, 6, 5, 4, 8]

heapq.heapreplace(someNumbers, 0) # извлекает 1 и вставляет 0
print(someNumbers) # ВЫВОД: [0, 3, 2, 7, 6, 5, 4, 8]

Метод heapreplace() отличается от heappushpop() только порядком вставки и извлечения:

• heapreplace() сперва извлекает корень кучи, а уже потом вставляет новый элемент.
• heappushpop() сперва вставляет новый элемент, а уже потом извлекает корень кучи.

Такое концептуально простое различие существенно влияет на результаты работы каждого из методов.

При использовании heapreplace():

• Всегда извлекается текущий корень кучи, независимо от того, больше он нового элемента или нет.
• Гарантированно удаляется старый корень и заменяется на другой элемент.

При использовании heappushpop():

• Возвращается либо старый корень кучи, либо только что вставленный элемент.
• Если вставленный элемент меньше старого корня, куча не меняется. Вместо этого удаляется только что вставленный элемент, словно его и не было в куче.

Таким образом, heapreplace() лучше использовать в тех случаях, когда необходимо безоговорочно удалить текущий корень и вставить новый элемент независимо от того, больше он или меньше старого корня.

Напротив, heappushpop() подойдет в тех случаях, когда точно неизвестно, больше новый элемент старого корня или, наоборот, меньше.

merge(список ... , ключ=None, инверсия=False) : итератор

Выполняет слияние нескольких разных куч в одну целую:

import heapq

firstNumbers = [1, 4, 9]
secondNumbers = [2, 6, 8]
thirdNumbers = [0, 3, 5, 7]

merged = heapq.merge(firstNumbers, secondNumbers, thirdNumbers) # создание генератора на основе списков

print(merged) # ВЫВОД: <generator object merge at 0x78b13064f4c0>
print(list(merged)) # ВЫВОД: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

При этом Python создает не сразу список всех элементов, а генератор — объект, который хранит состояние перебора и выдает элементы по одному при каждом запросе.

Во-первых, генератор не объединяет все списки сразу, а хранит внутренние указатели на текущие позиции каждого итерируемого объекта.

Во-вторых, генератор поэтапно выдает новые значения на основе текущих только в момент итерации через функцию next() или цикл for.

Результат работы функции можно инвертировать с помощью явно указанного параметра reverse=True:

import heapq

firstNumbers = [9, 4, 1]
secondNumbers = [8, 6, 2]
thirdNumbers = [7, 5, 3, 0]

merged = heapq.merge(firstNumbers, secondNumbers, thirdNumbers, reverse=True) # обратный порядок

print(list(merged)) # ВЫВОД: [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

Для использования параметра reverse=True исходные списки должны быть отсортированы в обратном порядке. В противном случае результат окажется некорректным — он не будет удовлетворять свойству кучи (max-heap):

import heapq

firstNumbers = [1, 4, 9]
secondNumbers = [2, 6, 8]
thirdNumbers = [0, 3, 5, 7]

merged = heapq.merge(firstNumbers, secondNumbers, thirdNumbers, reverse=True) # обратный порядок

print(list(merged)) # ВЫВОД: [2, 6, 8, 1, 4, 9, 0, 3, 5, 7]

Таким образом при reverse=False (установлен по умолчанию) функция ищет минимум среди элементов каждого списка и выдает их в порядке возрастания, а при reverse=True — ищет максимум и выдает элементы в порядке убывания.

С помощью дополнительного параметра key можно задать функцию сравнения:

import heapq

firstWords = ["kiwi", "apple", "banana"] # слова отсортированы на основе длины по возрастанию: 4, 5, 6
secondWords = ["fig", "pear", "grape"] # слова отсортированы на основе длины по возрастанию: 3, 4, 5

merged = heapq.merge(firstWords, secondWords, key=len)
print(list(merged)) # ВЫВОД: ['fig', 'kiwi', 'pear', 'apple', 'grape', 'banana']

#  Длины слов в новом списке по возрастанию: 3, 4, 4, 5, 5, 6

Помимо встроенных функций можно использовать пользовательские:

import heapq

firstWords = ["kiwi", "apple", "banana"]
secondWords = ["fig", "pear", "grape"]

def byLength(word):
	return len(word)

merged = heapq.merge(firstWords, secondWords, key=byLength)
print(list(merged)) # ВЫВОД: ['fig', 'kiwi', 'pear', 'apple', 'grape', 'banana']

Или лямбда-выражения:

import heapq

firstWords = ["kiwi", "apple", "banana"]
secondWords = ["fig", "pear", "grape"]

merged = heapq.merge(firstWords, secondWords, key=lambda w: len(w))
print(list(merged)) # ВЫВОД: ['fig', 'kiwi', 'pear', 'apple', 'grape', 'banana']

nlargest(количество, список, ключ=None) : список

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

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7]
maxNumbers = heapq.nlargest(3, someNumbers)
print(maxNumbers) # ВЫВОД: [8, 7, 6]

С помощью дополнительного параметра key можно задать функцию сравнения:

import heapq

somePeople = [
	{"name": "Саша", "age": 34},
	{"name": "Катя", "age": 28},
	{"name": "Дима", "age": 40},
	{"name": "Лена", "age": 22},
]

someOldPeople = heapq.nlargest(2, somePeople, key=lambda p: p["age"])

print(someOldPeople) # ВЫВОД: [{'name': 'Дима', 'age': 40}, {'name': 'Саша', 'age': 34}]

nsmallest(количество, список, ключ=None) : список

Возвращает список из указанного количество самых маленьких элементов кучи, упорядоченных по возрастания:

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7]
maxNumbers = heapq.nsmallest(3, someNumbers)
print(maxNumbers) # ВЫВОД: [1, 2, 3]

С помощью дополнительного параметра key можно задать функцию сравнения:

import heapq

somePeople = [
	{"name": "Саша", "age": 34},
	{"name": "Катя", "age": 28},
	{"name": "Дима", "age": 40},
	{"name": "Лена", "age": 22},
]

someYoungPeople = heapq.nsmallest(2, somePeople, key=lambda p: p["age"])

print(someYoungPeople) # ВЫВОД: [{'name': 'Лена', 'age': 22}, {'name': 'Катя', 'age': 28}]

Max-heap в Python через heapq: как обойти ограничение

Основная проблема модуля heapq — по умолчанию он реализует min-кучу. Чтобы перевернуть ее в max-кучу, есть два распространенных приема.

Способ 1. Отрицательные значения

С помощью отрицательных значений можно превратить min-кучу в max-кучи. Для этого необходимо всегда инвертировать элементы списка кучи:

• При вставке в min-кучу использовать -x вместо x.
• При извлечении из min-кучи использовать -heapq.heappop().

В этом случае по прежнему используется min-куча, однако механизм ее работы соответствует max-куче:

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7]
someNumbersInverted = []

for x in someNumbers:
	heapq.heappush(someNumbersInverted, -x)

heapq.heapify(someNumbersInverted)

print(someNumbersInverted) # ВЫВОД: [-8, -7, -5, -6, -3, -2, -4, -1]

Способ 2. Неофициальные функции

В язык программирования Python встроены неофициальные функции для работы с max-кучей — все они имеют нижнее подчеркивание в начале названия:

• _heapify_max()
• _heappop_max()
• _heapreplace_max()

С помощью этих функций можно выполнять самые базовые манипуляции на max-кучей:

import heapq

someNumbers = [8, 3, 5, 1, 6, 2, 4, 7]

heapq._heapify_max(someNumbers)
print(someNumbers) # ВЫВОД: [8, 7, 5, 3, 6, 2, 4, 1]

heapq._heappop_max(someNumbers)
print(someNumbers) # ВЫВОД: [7, 6, 5, 3, 1, 2, 4]

heapq._heapreplace_max(someNumbers, 9)
print(someNumbers) # ВЫВОД: [9, 6, 5, 3, 1, 2, 4]

Однако неофициальная реализация max-кучи имеет ограничения — в ней нет других привычных функций:

• _heappusр_max()
• _heappushpop_max()
• _merge_max()
• _nlargest_max()
• _nsmallest_max()

Пример использования heapq в задаче на Python

Для наглядной демонстрации работы кучи в Python можно рассмотреть реализацию простого планировщика задач:

import heapq
import itertools
import time

counter = itertools.count() # генератор уникальных идентификаторов, который позволяет при равных приоритетах извлекать задачи в порядке вставки задачи

# класс задачи

class Task:
	def __init__(self, name, priority):
		self.name = name
		self.priority = priority
		self.id = next(counter)
 
	def __repr__(self):
		return f"Task({self.name!r}, priority={self.priority})"

someTasks = [] # список для хранения задач

# функция для вставки задач

def schedule(task):
	heapq.heappush(someTasks, (-task.priority, task.id, task)) # задачи хранятся в списке кучи в виде кортежей, а их приоритет инвертируется для реализации max-кучи

# функция для обработки задач

def runNext():
	priority, id, task = heapq.heappop(someTasks) # извлечение задачи из кучи
	print(f"[{time.strftime('%H:%M:%S')}] Running:", task) # вывод задачи в консоль

# добавление нескольких задач

schedule(Task("Обработать запрос клиента", priority=50))
schedule(Task("Отправить отчёт", priority=20))
schedule(Task("Резервное копирование", priority=10))
schedule(Task("Критическое обновление", priority=100))

# выполнение задач в порядке убывания приоритета

while someTasks:
	runNext()
	time.sleep(0.5) # небольшая задержка для наглядности

Консольный вывод этого скрипта будет следующим:

[06:12:34] Running: Task('Критическое обновление', priority=100)
[06:12:34] Running: Task('Обработать запрос клиента', priority=50)
[06:12:35] Running: Task('Отправить отчёт', priority=20)
[06:12:35] Running: Task('Резервное копирование', priority=10)

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

Подготовили для вас выгодные тарифы на облачные серверы

Заключение

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

Работа с кучей сводится к нескольким функциям из модуля heapq — они автоматически поддерживают порядок без дополнительных усилий.

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

Простота применения в сочетании с универсальностью делает ее одним из основных инструментов для управления предельными значениями в наборах данных.