Skip to main content
`Ctrl`+`K`

# Source code for networkx.utils.mapped_queue

    """Priority queue class with updatable priorities."""

    import heapq

    __all__ = ["MappedQueue"]

    class _HeapElement:
        """This proxy class separates the heap element from its priority.

        The idea is that using a 2-tuple (priority, element) works
        for sorting, but not for dict lookup because priorities are
        often floating point values so round-off can mess up equality.

        So, we need inequalities to look at the priority (for sorting)
        and equality (and hash) to look at the element to enable
        updates to the priority.

        Unfortunately, this class can be tricky to work with if you forget that
        `__lt__` compares the priority while `__eq__` compares the element.
        In `greedy_modularity_communities()` the following code is
        used to check that two _HeapElements differ in either element or priority:

            if d_oldmax != row_max or d_oldmax.priority != row_max.priority:

        If the priorities are the same, this implementation uses the element
        as a tiebreaker. This provides compatibility with older systems that
        use tuples to combine priority and elements.
        """

        __slots__ = ["priority", "element", "_hash"]

        def __init__(self, priority, element):
            self.priority = priority
            self.element = element
            self._hash = hash(element)

        def __lt__(self, other):
            try:
                other_priority = other.priority
            except AttributeError:
                return self.priority < other
            # assume comparing to another _HeapElement
            if self.priority == other_priority:
                try:
                    return self.element < other.element
                except TypeError as err:
                    raise TypeError(
                        "Consider using a tuple, with a priority value that can be compared."
                    )
            return self.priority < other_priority

        def __gt__(self, other):
            try:
                other_priority = other.priority
            except AttributeError:
                return self.priority > other
            # assume comparing to another _HeapElement
            if self.priority == other_priority:
                try:
                    return self.element > other.element
                except TypeError as err:
                    raise TypeError(
                        "Consider using a tuple, with a priority value that can be compared."
                    )
            return self.priority > other_priority

        def __eq__(self, other):
            try:
                return self.element == other.element
            except AttributeError:
                return self.element == other

        def __hash__(self):
            return self._hash

        def __getitem__(self, indx):
            return self.priority if indx == 0 else self.element[indx - 1]

        def __iter__(self):
            yield self.priority
            try:
                yield from self.element
            except TypeError:
                yield self.element

        def __repr__(self):
            return f"_HeapElement({self.priority}, {self.element})"

    class MappedQueue:
        """The MappedQueue class implements a min-heap with removal and update-priority.

        The min heap uses heapq as well as custom written _siftup and _siftdown
        methods to allow the heap positions to be tracked by an additional dict
        keyed by element to position. The smallest element can be popped in O(1) time,
        new elements can be pushed in O(log n) time, and any element can be removed
        or updated in O(log n) time. The queue cannot contain duplicate elements
        and an attempt to push an element already in the queue will have no effect.

        MappedQueue complements the heapq package from the python standard
        library. While MappedQueue is designed for maximum compatibility with
        heapq, it adds element removal, lookup, and priority update.

        Parameters
        ----------
        data : dict or iterable

        Examples
        --------

        A `MappedQueue` can be created empty, or optionally, given a dictionary
        of initial elements and priorities.  The methods `push`, `pop`,
        `remove`, and `update` operate on the queue.

        >>> colors_nm = {"red": 665, "blue": 470, "green": 550}
        >>> q = MappedQueue(colors_nm)
        >>> q.remove("red")
        >>> q.update("green", "violet", 400)
        >>> q.push("indigo", 425)
        True
        >>> [q.pop().element for i in range(len(q.heap))]
        ['violet', 'indigo', 'blue']

        A `MappedQueue` can also be initialized with a list or other iterable. The priority is assumed
        to be the sort order of the items in the list.

        >>> q = MappedQueue([916, 50, 4609, 493, 237])
        >>> q.remove(493)
        >>> q.update(237, 1117)
        >>> [q.pop() for i in range(len(q.heap))]
        [50, 916, 1117, 4609]

        An exception is raised if the elements are not comparable.

        >>> q = MappedQueue([100, "a"])
        Traceback (most recent call last):
        ...
        TypeError: '<' not supported between instances of 'int' and 'str'

        To avoid the exception, use a dictionary to assign priorities to the elements.

        >>> q = MappedQueue({100: 0, "a": 1})

        References
        ----------
        .. [1] Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2001).
           Introduction to algorithms second edition.
        .. [2] Knuth, D. E. (1997). The art of computer programming (Vol. 3).
           Pearson Education.
        """

        def __init__(self, data=None):
            """Priority queue class with updatable priorities."""
            if data is None:
                self.heap = []
            elif isinstance(data, dict):
                self.heap = [_HeapElement(v, k) for k, v in data.items()]
            else:
                self.heap = list(data)
            self.position = {}
            self._heapify()

        def _heapify(self):
            """Restore heap invariant and recalculate map."""
            heapq.heapify(self.heap)
            self.position = {elt: pos for pos, elt in enumerate(self.heap)}
            if len(self.heap) != len(self.position):
                raise AssertionError("Heap contains duplicate elements")

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

        def push(self, elt, priority=None):
            """Add an element to the queue."""
            if priority is not None:
                elt = _HeapElement(priority, elt)
            # If element is already in queue, do nothing
            if elt in self.position:
                return False
            # Add element to heap and dict
            pos = len(self.heap)
            self.heap.append(elt)
            self.position[elt] = pos
            # Restore invariant by sifting down
            self._siftdown(0, pos)
            return True

        def pop(self):
            """Remove and return the smallest element in the queue."""
            # Remove smallest element
            elt = self.heap[0]
            del self.position[elt]
            # If elt is last item, remove and return
            if len(self.heap) == 1:
                self.heap.pop()
                return elt
            # Replace root with last element
            last = self.heap.pop()
            self.heap[0] = last
            self.position[last] = 0
            # Restore invariant by sifting up
            self._siftup(0)
            # Return smallest element
            return elt

        def update(self, elt, new, priority=None):
            """Replace an element in the queue with a new one."""
            if priority is not None:
                new = _HeapElement(priority, new)
            # Replace
            pos = self.position[elt]
            self.heap[pos] = new
            del self.position[elt]
            self.position[new] = pos
            # Restore invariant by sifting up
            self._siftup(pos)

        def remove(self, elt):
            """Remove an element from the queue."""
            # Find and remove element
            try:
                pos = self.position[elt]
                del self.position[elt]
            except KeyError:
                # Not in queue
                raise
            # If elt is last item, remove and return
            if pos == len(self.heap) - 1:
                self.heap.pop()
                return
            # Replace elt with last element
            last = self.heap.pop()
            self.heap[pos] = last
            self.position[last] = pos
            # Restore invariant by sifting up
            self._siftup(pos)

        def _siftup(self, pos):
            """Move smaller child up until hitting a leaf.

            Built to mimic code for heapq._siftup
            only updating position dict too.
            """
            heap, position = self.heap, self.position
            end_pos = len(heap)
            startpos = pos
            newitem = heap[pos]
            # Shift up the smaller child until hitting a leaf
            child_pos = (pos << 1) + 1  # start with leftmost child position
            while child_pos < end_pos:
                # Set child_pos to index of smaller child.
                child = heap[child_pos]
                right_pos = child_pos + 1
                if right_pos < end_pos:
                    right = heap[right_pos]
                    if not child < right:
                        child = right
                        child_pos = right_pos
                # Move the smaller child up.
                heap[pos] = child
                position[child] = pos
                pos = child_pos
                child_pos = (pos << 1) + 1
            # pos is a leaf position. Put newitem there, and bubble it up
            # to its final resting place (by sifting its parents down).
            while pos > 0:
                parent_pos = (pos - 1) >> 1
                parent = heap[parent_pos]
                if not newitem < parent:
                    break
                heap[pos] = parent
                position[parent] = pos
                pos = parent_pos
            heap[pos] = newitem
            position[newitem] = pos

        def _siftdown(self, start_pos, pos):
            """Restore invariant. keep swapping with parent until smaller.

            Built to mimic code for heapq._siftdown
            only updating position dict too.
            """
            heap, position = self.heap, self.position
            newitem = heap[pos]
            # Follow the path to the root, moving parents down until finding a place
            # newitem fits.
            while pos > start_pos:
                parent_pos = (pos - 1) >> 1
                parent = heap[parent_pos]
                if not newitem < parent:
                    break
                heap[pos] = parent
                position[parent] = pos
                pos = parent_pos
            heap[pos] = newitem
            position[newitem] = pos