I am trying to implement a priority queue class using the heapq
library, which requires that items inserted into the queue are comparable with each other.
I would like for the constructor of the queue class to take in a comparator function which, when passed two items, will return a number less than, equal to, or greater than 0 to indicate the relationship between the priorities of the two items. Using this comparator function, I somehow need to create objects to be inserted into the queue which can be directly compared with each other. Here is what I have right now:
from functools import total_ordering
from heapq import heappush, heappop
class PriorityQueue:
__REMOVED = "<removed-task>"
def __init__(self, comparator):
self._pq = [] # list of entries arranged in a heap
self._entry_finder = {} # mapping of tasks to entries
@total_ordering
class PrioritizedItem:
def __init__(self, task):
self._task = task
@property
def task(self):
return self._task
@task.setter
def task(self, task):
self._task = task
def __eq__(self, other):
return comparator(self._task, other._task) == 0
def __lt__(self, other):
return comparator(self._task, other._task) < 0
self._item_class = PrioritizedItem
def clear(self):
"Removes all of the elements from this priority queue. The queue will be empty after this call returns."
self._pq.clear()
self._entry_finder.clear()
def remove(self, task):
"Mark an existing task as REMOVED. Raise KeyError if not found."
entry = self._entry_finder.pop(task)
entry.task = self.__REMOVED
def add(self, task):
"Add a new task or update the priority of an existing task"
if task in self._entry_finder:
self.remove(task)
entry = this._item_class(task)
self._entry_finder[task] = entry
heappush(self._pq, entry)
def pop(self):
"Remove and return the lowest priority task. Raise KeyError if empty."
while self._pq:
task = heappop(self._pq).task
if task is not self.__REMOVED:
del self._entry_finder[task]
return task
raise KeyError("pop from an empty priority queue")
I didn't want to make the PrioritizedItem
class a standalone class outside of PriorityQueue
since then I would need to pass the comparator function into the constructor for each instance of the class. It doesn't really make sense to me to define an ordered class if every instance of the class could potentially contain a different comparator. (For instance, if two instances x
and y
were created with different comparator functions, then we could have x == y
but y != x
).
On the other hand, I'm not sure if my current approach is great either since I'm defining a class inside the constructor of the queue and assigning the class to an instance variable, which just feels weird.
Is there any way to improve this code? Thanks!
Aucun commentaire:
Enregistrer un commentaire