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’ полезно иметь под рукой. :-)

Сноски