heapq — Алгоритм очереди с кучей

Основные примеры

Сортировку кучей можно реализовать, добавив все значения в кучу, а затем по одному извлекая наименьшие значения:

>>> def heapsort(iterable):
...     h = []
...     for value in iterable:
...         heappush(h, value)
...     return [heappop(h) for i in range(len(h))]
...
>>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Это похоже на sorted(iterable), но, в отличие от sorted(), эта реализация не является стабильной.

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

>>> h = []
>>> heappush(h, (5, 'write code'))
>>> heappush(h, (7, 'release product'))
>>> heappush(h, (1, 'write spec'))
>>> heappush(h, (3, 'create tests'))
>>> heappop(h)
(1, 'write spec')

Примечания по реализации очереди с приоритетами

Очередь с приоритетами — распространённый пример использования кучи, который ставит перед разработчиком несколько задач:

• Стабильность сортировки: как добиться того, чтобы две задачи с одинаковыми приоритетами возвращались в порядке их добавления?
• Сравнение кортежей нарушается для пар (приоритет, задача), если приоритеты равны, а для задач не определён порядок сравнения по умолчанию.
• Если приоритет задачи меняется, как переместить её на новое место в куче?
• Или, если ожидающую выполнения задачу нужно удалить, как найти её и убрать из очереди?

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

Ещё одно решение проблемы несравнимых задач — создать класс-обёртку, который игнорирует элемент задачи и сравнивает только поле приоритета:

from dataclasses import dataclass, field
from typing import Any

@dataclass(order=True)
class PrioritizedItem:
    priority: int
    item: Any=field(compare=False)

Оставшиеся задачи связаны с поиском ожидающей выполнения задачи и изменением её приоритета или полным удалением. Задачу можно найти с помощью словаря, указывающего на запись в очереди.

Удалить запись или изменить её приоритет сложнее, поскольку это нарушит инварианты структуры кучи. Возможное решение — пометить запись как удалённую и добавить новую запись с изменённым приоритетом:

pq = []                         # list of entries arranged in a heap
entry_finder = {}               # mapping of tasks to entries
REMOVED = '<removed-task>'      # placeholder for a removed task
counter = itertools.count()     # unique sequence count

def add_task(task, priority=0):
    'Add a new task or update the priority of an existing task'
    if task in entry_finder:
        remove_task(task)
    count = next(counter)
    entry = [priority, count, task]
    entry_finder[task] = entry
    heappush(pq, entry)

def remove_task(task):
    'Mark an existing task as REMOVED.  Raise KeyError if not found.'
    entry = entry_finder.pop(task)
    entry[-1] = REMOVED

def pop_task():
    'Remove and return the lowest priority task. Raise KeyError if empty.'
    while pq:
        priority, count, task = heappop(pq)
        if task is not REMOVED:
            del entry_finder[task]
            return task
    raise KeyError('pop from an empty priority queue')

Теория

Кучи — это массивы, для которых a[k] <= a[2*k+1] и a[k] <= a[2*k+2] для всех k, если считать элементы с нуля. Для сравнения несуществующие элементы считаются бесконечными. Интересное свойство кучи заключается в том, что a[0] всегда является её наименьшим элементом.

Описанный выше необычный инвариант предназначен для эффективного представления турнира в памяти. Числа внизу — это k, а не a[k]:

                               0

              1                                 2

      3               4                5               6

  7       8       9       10      11      12      13      14

15 16   17 18   19 20   21 22   23 24   25 26   27 28   29 30

В приведённом выше дереве каждая ячейка k находится над 2*k+1 и 2*k+2. В обычном спортивном бинарном турнире каждая ячейка представляет победителя двух ячеек, над которыми она находится, и победителя можно проследить вниз по дереву, чтобы увидеть всех его соперников. Однако во многих компьютерных приложениях таких турниров историю победителя отслеживать не нужно. Чтобы эффективнее использовать память, при продвижении победителя мы стараемся заменить его чем-нибудь другим на нижнем уровне. В результате ячейка и две ячейки под ней содержат три разных элемента, но верхняя ячейка «побеждает» две нижние.

Если этот инвариант кучи постоянно соблюдается, индекс 0 явно соответствует общему победителю. Простейший алгоритмический способ удалить его и найти «следующего» победителя — переместить какого-нибудь проигравшего (скажем, ячейку 30 на схеме выше) в позицию 0, а затем просеивать новый элемент из позиции 0 вниз по дереву, меняя значения местами, пока инвариант не будет восстановлен. Очевидно, что эта операция требует логарифмического времени относительно общего количества элементов в дереве. Повторяя её для всех элементов, получим сортировку со сложностью O(n log n).

Приятная особенность этой сортировки — возможность эффективно добавлять новые элементы прямо во время сортировки, если добавляемые элементы не «лучше» последнего извлечённого элемента из позиции 0. Это особенно полезно при моделировании, где дерево хранит все поступающие события, а условие «победы» означает наиболее раннее запланированное время. Когда событие планирует другие события для выполнения, они назначаются на будущее, поэтому их легко добавить в кучу. Таким образом, куча — хорошая структура для реализации планировщиков (именно её я использовал в своём MIDI-секвенсоре :-).

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

Кучи также очень полезны при больших внешних сортировках. Скорее всего, вы знаете, что большая сортировка подразумевает создание «серий» (предварительно отсортированных последовательностей, размер которых обычно связан с объёмом памяти процессора), за которыми следуют проходы слияния этих серий. Такое слияние часто организуется весьма изобретательно [1]. Очень важно, чтобы начальная сортировка создавала как можно более длинные серии. Турниры хорошо подходят для этого. Если, используя всю доступную память для хранения турнира, заменять и просеивать элементы, которые подходят для текущей серии, то для случайных входных данных длина получаемых серий будет вдвое больше объёма памяти, а для входных данных, уже отчасти упорядоченных, результат будет ещё лучше.

Кроме того, если вывести на диск элемент с индексом 0 и получить входное значение, которое не подходит для текущего турнира (поскольку это значение «побеждает» последнее выведенное значение), его нельзя добавить в кучу, поэтому размер кучи уменьшается. Освободившуюся память можно сразу же с пользой применить для постепенного построения второй кучи, которая будет расти ровно с той же скоростью, с какой уменьшается первая. Когда первая куча полностью исчезнет, можно переключиться на вторую и начать новую серию. Умно и весьма эффективно!

Словом, кучи — полезные структуры данных, о которых стоит знать. Я использую их в нескольких приложениях и считаю, что модуль «heap» всегда пригодится. :-)

Сноски