How can I implement decrease-key functionality in Python's heapq?

Viewed 23012

I know it is possible to realize decrease-key functionality in O(log n) but I don't know how?

7 Answers

It might be unnecessary to have the decrease_key function (albeit it's nice to have it).

You can just push your (priority, item) into the priority queue anyway, and use a set to check whether you have seen it. For example:

pq = []  # heapq is a min heap
seen = set()
heappush(pq, (2, "item1"))
heappush(pq, (3, "item2"))
heappush(pq, (1, "item3"))
heappush(pq, (4, "item4"))
heappush(pq, (2, "item2"))

while pq:
    p, item = heappop(pq)
    if item not in seen:
        seen.add(item)
        print(item, p)
    else:
        print(item, "is already handled with a higher priority!")

The output is:

item3 1
item1 2
item2 2
item2 is already handled with a higher priority!
item4 4

This functionality is also missing from C++ and Java standard library priority queues. The standard workaround is to push a new key-value pair and either implicitly or explicitly mark the original key-value pair as invalid.

See How to update elements within a heap? (priority queue) and Why does Dijkstra's algorithm use decrease-key? (the conclusion there is that not having decrease-key won't significantly impact runtime in theory and in practice; in particular see the paper https://www3.cs.stonybrook.edu/~rezaul/papers/TR-07-54.pdf)

Priority Queue Simple Implementation with Min heap with unique keys. This implementation updates the priority of the key during push() operation if key already exists in priority queue.

import heapq

class HeapPQ:
    """
    Only hashable key type is supported
    """

    def __init__(self):
        self.pq = []
        self.pq_set = set()
    
    def push(self, priority, key):
        if key not in self.pq_set:
            heapq.heappush(self.pq, (priority, key))
            self.pq_set.add(key)
        else:
            index = list(map(lambda x:x[1], self.pq)).index(key)
            self.pq[index] = (priority, key)
            heapq.heapify(self.pq)
    
    def pop(self):
        priority, key = heapq.heappop(self.pq)
        self.pq_set.remove(key)
        return priority, key
    
    def empty(self) -> bool:
        return len(self.pq) == 0

Example Usage:

pq = HeapPQ()

pq.push(5, "A")
pq.push(3, "B")
pq.push(1, "A")

while not pq.empty():
    print(pq.pop())
Related