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

Источник: Lib/heapq.py

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

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

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

Приведенный ниже API отличается от алгоритмов кучи, описанных в учебниках, двумя аспектами: (a) Мы используем индексацию на основе нуля. Это делает связь между индексом узла и индексами его дочерних элементов несколько менее очевидной, но более подходящей, поскольку в Python используется нулевая индексация. (b) Наш метод pop возвращает наименьший элемент, а не наибольший (в учебниках это называется «min heap»; «max heap» более распространен в текстах из-за его пригодности для сортировки на месте).

Эти два параметра позволяют рассматривать кучу как обычный список Python без сюрпризов: heap[0] - наименьший элемент, а heap.sort() сохраняет инвариант кучи!

Чтобы создать кучу, используйте список, инициализированный значением [], или вы можете преобразовать заполненный список в кучу с помощью функции heapify().

Предусмотрены следующие функции:

heapq.heappush(heap, item)

Переместите значение item в кучу, сохраняя инвариант кучи.

heapq.heappop(heap)

Извлеките и верните наименьший элемент из кучи, сохраняя инвариант кучи. Если куча пуста, вызывается IndexError. Чтобы получить доступ к наименьшему элементу, не вынимая его, используйте heap[0].

heapq.heappushpop(heap, item)

Переместите item в кучу, затем выньте и верните наименьший элемент из кучи. Комбинированное действие работает эффективнее, чем heappush() с последующим отдельным вызовом heappop().

heapq.heapify(x)

Преобразование списка x в кучу, in-place, за линейное время.

heapq.heapreplace(heap, item)

Вытащить и вернуть наименьший элемент из кучи, а также протолкнуть новый элемент. Размер кучи не изменяется. Если куча пуста, возвращается значение IndexError.

Эта одношаговая операция более эффективна, чем heappop() с последующим heappush(), и может быть более подходящей при использовании кучи фиксированного размера. Комбинация pop/push всегда возвращает элемент из кучи и заменяет его на item.

Возвращаемое значение может быть больше, чем добавленный элемент. Если это нежелательно, используйте вместо этого heappushpop(). Его комбинация push/pop возвращает меньшее из двух значений, оставляя большее значение в куче.

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

heapq.merge(*iterables, key=None, reverse=False)

Объединяет несколько отсортированных входных данных в один отсортированный выход (например, объединяет записи с временными метками из нескольких файлов журнала). Возвращает iterator над отсортированными значениями.

Аналогично sorted(itertools.chain(*iterables)), но возвращает итерируемый поток, не забирает данные в память все сразу и предполагает, что каждый из входных потоков уже отсортирован (от меньшего к большему).

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

key задает key function одного аргумента, который используется для извлечения ключа сравнения из каждого входного элемента. По умолчанию используется значение None (сравнение элементов напрямую).

reverse - булево значение. Если установлено значение True, то входные элементы объединяются так, как если бы каждое сравнение было обратным. Чтобы добиться поведения, аналогичного sorted(itertools.chain(*iterables), reverse=True), все итеративные элементы должны быть отсортированы от наибольшего к наименьшему.

Изменено в версии 3.5: Добавлены необязательные параметры key и reverse.

heapq.nlargest(n, iterable, key=None)

Возвращает список с n наибольшими элементами из набора данных, заданного iterable. key, если указано, задает функцию с одним аргументом, которая используется для извлечения ключа сравнения из каждого элемента в iterable (например, key=str.lower). Эквивалентно: sorted(iterable, key=key, reverse=True)[:n].

heapq.nsmallest(n, iterable, key=None)

Возвращает список с n наименьшими элементами из набора данных, заданного iterable. key, если указано, задает функцию с одним аргументом, которая используется для извлечения ключа сравнения из каждого элемента в iterable (например, key=str.lower). Эквивалентно: sorted(iterable, key=key)[:n].

Две последние функции лучше всего работают при небольших значениях n. Для больших значений эффективнее использовать функцию sorted(). Кроме того, при значении n==1 эффективнее использовать встроенные функции min() и max(). Если требуется многократное использование этих функций, подумайте о том, чтобы превратить итерируемую переменную в настоящую кучу.

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

Можно реализовать heapsort, поместив все значения в кучу, а затем выгружая наименьшие значения по одному за раз:

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

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

Для кучи обычно используется priority queue, и это создает несколько проблем при реализации:

  • Стабильность сортировки: как сделать так, чтобы две задачи с одинаковыми приоритетами возвращались в том порядке, в котором они были изначально добавлены?

  • Сравнение кортежей прерывается для пар (приоритет, задача), если приоритеты равны и задачи не имеют порядка сравнения по умолчанию.

  • Если приоритет задачи меняется, как переместить ее на новую позицию в куче?

  • Или если отложенную задачу нужно удалить, как найти ее и удалить из очереди?

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

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

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

Одним словом, кучи - это полезные структуры памяти, которые нужно знать. Я использую их в нескольких приложениях и считаю, что полезно иметь под рукой модуль „heap“ :-)

Сноски