import heapq from functools import total_ordering @total_ordering class SpecialSorted: def __init__(self, element, value): self.element = element self.value = value def __eq__(self, other): return self.value == other.value def __ne__(self, other): return self.value != other.value def __lt__(self, other): return self.value < other.value class PriorityQueue: def __init__(self, sortkey = lambda x : x): self.content = [] self.sortkey = sortkey def add(self, item): heapq.heappush(self.content, SpecialSorted(item, self.sortkey(item))) def peek(self): return self.content[0].element if self.content else None def poll(self): return heapq.heappop(self.content).element if len(self.content) > 0 else None def is_empty(self): return len(self.content) == 0 def __str__(self): return str(heapq.nsmallest(len(self.content), [item.element for item in self.content]))