Heapq

Use to create priority queue. The reason we use heapq instead of python PriorityQueue is because we can peek the first element and access to more internal apis.

Although heapq is not thread-safe. So when dealing with multi-threaded environment, we might need to use PriorityQueue.

import heapq

myHeap = [] 
heapq.heappush(myHeap, 5)
heapq.heappush(myHeap, 1)
heapq.heappush(myHeap, 8)
heapq.heappush(myHeap, 2)
heapq.heappush(myHeap, 0)
heapq.heappush(myHeap, 9)

print(myHeap) # [0, 1, 8, 5, 2, 9]

Note: only the first number here (0) is correct, the rest could be in any order.

So now, we have created a heap with smallest element first. To sort this heap, we can call heappop() as we go:

def iterate(heap: List[int]) -> int:
    while heap:
        yield heapq.heappop(heap)

print([sortedItem for sortedItem in iterate(myHeap)])

Using with tuple

We can store the tuple as the format (priority, data) to use with heapq. For example:

def iterate(heap: List[int]) -> int:
    while heap:
        yield heapq.heappop(heap)[1]

if __name__ == "__main__":
    myHeap = [] 
    
    heapq.heappush(myHeap, (5, "task5"))
    heapq.heappush(myHeap, (1, "task1"))
    heapq.heappush(myHeap, (8, "task8"))
    heapq.heappush(myHeap, (2, "task2"))
    heapq.heappush(myHeap, (0, "task0"))
    heapq.heappush(myHeap, (9, "task9"))
    
    
    print([sortedItem for sortedItem in iterate(myHeap)])

['task0', 'task1', 'task2', 'task5', 'task8', 'task9']

[!caution]
If the first element equals eachother, heapq will then start comparing the second element. As a result if you have something like heappush(myHeap, (5, MyClass) it will fail!

The way to fix this is to also pass in an id() so that when we do heappush it also checks for id as well:

heapq.heapush(myHeap, (5, id(MyClass), MyClass))

^eb63dd

heappushpop(heap, item)

Sometimes we want to push something and pop out, it's more effective to do heappushpop() than heappush() and then heappop().

For example:

def smallestNElements(n, array: List[int]) -> List[int]:
    if n >= len(array): raise Exception()

    heap = []
    for index, num in enumerate(array):
        if index >= n:
            heapq.heappushpop(heap, -num)
        else:
            heapq.heappush(heap, -num)

    while heap:
        item = -heapq.heappop(heap)
        yield item

print([item for item in sorted(smallestNElements(3, [5,3,8,1,10]))])

[1, 3, 5]

nsmallest(n, iterable)

Python implementation to find n-th smallest element

print(heapq.nsmallest(3, [5,3,8,1,10])) # [1, 3, 5]

nlargest(n, iterable)

Similarly, for n-th largest element

print(heapq.nlargest(3, [5,3,8,1,10])) # [10,8,5]