Zgodnie z dokumentacją heapq , sposobem dostosowania kolejności sterty jest ustawienie każdego elementu na stercie jako krotki, przy czym pierwszy element krotki to taki, który akceptuje zwykłe porównania w Pythonie.
Funkcje w module heapq są nieco kłopotliwe (ponieważ nie są zorientowane obiektowo) i zawsze wymagają, aby nasz obiekt sterty (lista sterty) był jawnie przekazany jako pierwszy parametr. Możemy upiec dwie pieczenie na jednym ogniu, tworząc bardzo prostą klasę opakowującą, która pozwoli nam określić key
funkcję i przedstawić stertę jako obiekt.
Poniższa klasa przechowuje wewnętrzną listę, w której każdy element jest krotką, której pierwszym składnikiem jest klucz, obliczany w czasie wstawiania elementu za pomocą key
parametru, przekazywany w instancji Heap:
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]
(Dodatkową self.index
częścią jest uniknięcie kolizji, gdy oszacowana wartość klucza jest remisem, a przechowywana wartość nie jest bezpośrednio porównywalna - w przeciwnym razie heapq może zawieść z TypeError)