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')

Другие применения

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

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

def running_median(iterable):
    "Yields the cumulative median of values seen so far."

    lo = []  # max-heap
    hi = []  # min-heap (same size as or one smaller than lo)

    for x in iterable:
        if len(lo) == len(hi):
            heappush_max(lo, heappushpop(hi, x))
            yield lo[0]
        else:
            heappush(hi, heappushpop_max(lo, x))
            yield (lo[0] + hi[0]) / 2

Например:

>>> list(running_median([5.0, 9.0, 4.0, 12.0, 8.0, 9.0]))
[5.0, 7.0, 5.0, 7.0, 8.0, 8.5]

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

Очередь с приоритетами — распространённое применение кучи, которое сопряжено с несколькими сложностями реализации:

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

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

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

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, при этом элементы отсчитываются от 0. Для сравнения несуществующие элементы считаются бесконечными. Примечательное свойство кучи заключается в том, что 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 и получить входное значение, которое не может войти в текущий турнир (поскольку это значение «побеждает» последнее выведенное значение), оно не поместится в кучу, и размер кучи уменьшится. Освободившуюся память можно сразу же эффективно использовать для постепенного построения второй кучи, которая растёт ровно с той же скоростью, с какой уменьшается первая. Когда первая куча полностью опустеет, можно переключиться на вторую и начать новую серию. Умно и довольно эффективно!

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

Сноски