Согласно документации heapq , способ настроить порядок кучи - сделать так, чтобы каждый элемент в куче был кортежем, причем первый элемент кортежа был тем, который принимает обычные сравнения Python.
Функции в модуле heapq немного громоздки (поскольку они не объектно-ориентированы) и всегда требуют, чтобы наш объект кучи (список с динамической памятью) явно передавался в качестве первого параметра. Мы можем убить двух зайцев одним выстрелом, создав очень простой класс-оболочку, который позволит нам указать key
функцию и представить кучу как объект.
В приведенном ниже классе хранится внутренний список, в котором каждый элемент является кортежем, первый член которого является ключом, вычисляемым во время вставки элемента с использованием key
параметра, переданного при создании экземпляра кучи:
import heapq
class MyHeap(object):
def __init__(self, initial=None, key=lambda x:x):
self.key = key
self.index = 0
if initial:
self._data = [(key(item), i, item) for i, item in enumerate(initial)]
self.index = len(self._data)
heapq.heapify(self._data)
else:
self._data = []
def push(self, item):
heapq.heappush(self._data, (self.key(item), self.index, item))
self.index += 1
def pop(self):
return heapq.heappop(self._data)[2]
(Дополнительная self.index
часть состоит в том, чтобы избежать конфликтов, когда оцененное значение ключа является ничьей, а сохраненное значение не может быть напрямую сопоставлено - иначе heapq может завершиться ошибкой с TypeError)