7
\$\begingroup\$

Here is my attempt to implement a minimum priority queue class in Python using the heapq module. I plan to use it in graph search algorithms so the values associated with each key are update-able.

Note: Any obsolete values related to an updated key are kept until they are on the top of the queue, at which time they are ignored.

from heapq import heappush, heappop, heapify


class PriorityQueueUpdateable(object):
    """
    Updateable Priority Queue implemented with heapq module and standard python dictionary
    """
    def __init__(self):
        self._heap = []
        self._dict = dict()
        heapify(self._heap)

    def _clear_heap(self):
        """
        Removes obsolete entries from heap
        """
        value, key = self._heap[0]
        while (key not in self._dict) or (self._dict[key] != value):
            heappop(self._heap)
            if not self._heap:
                break
            value, key = self._heap[0]

    def pop(self):
        if not self:
            raise IndexError("Queue is empty")

        self._clear_heap()

        value, key = heappop(self._heap)
        del self._dict[key]

        self._clear_heap()

        return key, value

    def peek(self):
        if not self:
            raise IndexError("Queue is empty")
        self._clear_heap()
        value, key = self._heap[0]
        return key, value

    def __contains__(self, key):
        return key in self._dict

    def __getitem__(self, key):
        return self._dict[key]

    def __len__(self):
        return len(self._dict)

    def __setitem__(self, key, value):
        self._dict[key] = value
        heappush(self._heap, (value, key))

The code passes this simple test

pq = PriorityQueueUpdateable()
pq[1]= 0.1
pq[1]= 2
pq[1]= 0.2
pq[2] = 6
pq[2] = 0.1
pq[2]= 0.001
assert pq[1] == 0.2
assert pq[2] == 0.001

assert 2 in pq
assert 1 in pq

assert len(pq) == 2
assert pq.pop() == (2, 0.001)
assert pq.pop() == (1, 0.2)

I would appreciate any feedback on improving my code. Thanks.

\$\endgroup\$
6
\$\begingroup\$

1. Bug

You can't pop an item from the queue if it's the only item:

>>> q = PriorityQueueUpdateable()
>>> q[1] = 1
>>> q.pop()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "cr163560.py", line 41, in pop
    self._clear_heap()
  File "cr163560.py", line 25, in _clear_heap
    value, key = self._heap[0]
IndexError: list index out of range

The problem is in the pop method: it calls _clear_heap after popping the item (as well as before). But at this point there might be no items left in the queue. In your simple test, you add both keys to the queue multiple times, so that when you pop an item it's never the last item (because there are more items with the same key remaining in the queue).

2. Review

  1. The Python library documentation has a section "Priority Queue Implementation Notes" which gives some advice on implementing a priority queue using a heap. This is well worth reading.

  2. The docstring for the class doesn't give much of a clue as to how to use it. Docstrings should be written from the user's point of view, and shouldn't include implementation details (like "implemented with heapq module") — these are only needed by developers and so can be mentioned in comments.

  3. There are no docstrings for the __init__, pop and peek methods.

  4. It would be convenient if the __init__ method accepted an optional iterable of (key, priority) pairs. This would make it easier to create priority queues.

  5. In __init__, there's no need to call heapq.heapify — this does nothing on an empty heap.

  6. Using __setitem__ to add an item to the queue seems strange — normally one would expect a method adding something to a queue to be named something like push, append, or put. (Compare heapq.heappush, collections.deque.append, queue.Queue.put.)

  7. The second argument to __setitem__ is named value, but this is misleading — actually it's the priority of the key.

  8. Priorities work in reverse — if you want an item to have a higher priority then you have to give it a lower value (this happens because Python's heaps are min-heaps). It would probably suit more applications if priorities worked in the usual way. (You'd have to negate the priority values to make higher values mean higher priorities.)

  9. You give a test case, which is great! It would be convenient to turn this into a unit test (perhaps using the unittest module) so that it can be automatically run.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.