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