Очередь с приоритетом в Python: heapq или PriorityQueue, в чём разница и как пользоваться?
Нужна очередь с приоритетом: всегда доставать элемент с наименьшим ключом. Нашёл два варианта — модуль heapq и класс queue.PriorityQueue. Не пойму, что выбрать и как вообще класть элементы с приоритетом (а не просто числа).
И ещё: heapq — это min-куча, а мне иногда нужна max. Как развернуть?
2 ответа
heapq — это не контейнер, а набор функций, которые превращают ОБЫЧНЫЙ список в min-кучу (двоичную). Самые ходовые:
import heapq
h = []
heapq.heappush(h, 5)
heapq.heappush(h, 1)
heapq.heappush(h, 3)
print(heapq.heappop(h)) # 1 — всегда наименьший
print(heapq.heappop(h)) # 3
# превратить готовый список в кучу за O(n):
nums = [9, 2, 7, 1]
heapq.heapify(nums)
print(heapq.heappop(nums)) # 1
Чтобы хранить элементы с приоритетом отдельно от данных, кладут кортежи (приоритет, данные) — куча сравнивает по первому элементу:
pq = []
heapq.heappush(pq, (2, "помыть посуду"))
heapq.heappush(pq, (1, "потушить пожар"))
heapq.heappush(pq, (3, "полить цветы"))
print(heapq.heappop(pq)) # (1, 'потушить пожар')
Ловушка: если приоритеты совпадут, Python начнёт сравнивать ВТОРОЙ элемент кортежа, и если данные несравнимы (например, объекты) — упадёт TypeError. Лечится добавлением уникального счётчика-тай-брейкера: (приоритет, count, данные).
Про max-кучу: heapq всегда min. Хитрость — класть отрицательный приоритет:
heapq.heappush(h, -value) # потом heappop вернёт -max, берёшь -результат
Когда что: heapq — быстрый, простой, для однопоточного кода и алгоритмов (Дейкстра, k наибольших через heapq.nlargest). queue.PriorityQueue — обёртка над тем же heapq, но потокобезопасная (с замком внутри) и с блокирующим get(). Бери её, только если очередь реально шарится между потоками, иначе она просто медленнее из-за лишних локов.
Готовые шорткаты из heapq, про которые часто забывают: heapq.nlargest(k, data) и heapq.nsmallest(k, data) — берут k наибольших/наименьших без ручного построения кучи. А heapq.heappushpop(h, x) и heapq.heapreplace(h, x) делают push+pop за один проход, эффективнее двух отдельных вызовов — удобно для поддержания «топ-k».