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

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

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

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

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

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

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

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

Примечания