BucketQueue.py

"""BucketQueue.py

A simple priority queue with integer priorities. Use like:
    Q = BucketQueue()
    Q[item] = priority
    del Q[item]
    for item in Q: ...

The time to find and return each item in the queue is proportional
to the difference (current priority - previous priority), or O(1) if
this difference is non-positive. In particular when the minimum priority
is non-negative and non-decreasing (as in Dijkstra's algorithm) or when
its decreases are bounded (as in graph degeneracy) the total time for a
sequence of operations is proportional to the number of operations plus
the maximum priority.

D. Eppstein, July 2016.
"""

class BucketQueue:
    def __init__(self):
        """Create a new empty integer priority queue."""
        self._D = {}        # map from items to priorities
        self._Q = {}        # map from priorities to buckets
        self._N = None      # lower bound on min priority

    def __getitem__(self,item):
        """Look up the priority of an item."""
        return self._D[item]

    def __delitem__(self,item):
        """Remove an item from the priority queue."""
        priority = self._D[item]
        del self._D[item]               # remove from map of items => priorities
        self._Q[priority].remove(item)  # remove from bucket
        if not self._Q[priority]:
            del self._Q[priority]       # remove empty bucket

    def __setitem__(self,item,priority):
        """Add an element to the priority queue with the given priority."""
        if not isinstance(priority,int):
            raise TypeError("Priority must be an integer")
        if item in self._D:
            del self[item]
        self._D[item] = priority        # add to map of items => priorities
        if not self._Q or priority < self._N:
            self._N = priority          # update priority lower bound
        if priority not in self._Q:
            self._Q[priority] = set()   # make new bucket if necessary
        self._Q[priority].add(item)     # and add to bucket

    def __iter__(self):
        """Repeatedly find and remove the min-priority item from the queue.
        It is ok for the queue to be modified between iterations."""
        while self._Q:
            while self._N not in self._Q:
                self._N += 1
            x = next(iter(self._Q[self._N]))    # arbitrary item in 1st bucket
            del self[x]
            yield x

    def items(self):
        """Variant iterator that generates (item,priority) pairs.
        We rely on the fact that the usual __iter__ always
        leaves self._N equal to the priority."""
        for x in iter(self):
            yield x,self._N

    def __contains__(self,item):
        """Container class membership test."""
        return item in self._D

    def __len__(self):
        """Container class length."""
        return len(self._D)