Python Heapq Module
The heapq module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm. Heaps are binary trees for which every parent node has a value less than or equal to any of its children. This implementation uses a min-heap.
Importing the Module
import heapqBasic Operations
Creating a Heap
You can use a list as a heap. To convert a populated list into a heap, use heapify().
import heapq
h = [3, 1, 4, 1, 5, 9, 2, 6]
heapq.heapify(h)
print(h)
# Output: [1, 1, 2, 3, 5, 9, 4, 6] (Order may vary, but h[0] is min)Pushing Items
Use heappush() to add an item to the heap while maintaining the heap invariant.
heapq.heappush(h, 0)
print(h)
# Output: [0, 1, 2, 1, 5, 9, 4, 6, 3]Popping Items
Use heappop() to remove and return the smallest item from the heap.
smallest = heapq.heappop(h)
print(smallest) # Output: 0Efficient Push/Pop
heappushpop()
Pushes an item on the heap and then pops and returns the smallest item. This is more efficient than a separate push followed by a pop.
result = heapq.heappushpop(h, 7)
print(result)heapreplace()
Pops and returns the smallest item from the heap, and then pushes the new item. The heap size does not change. This is more efficient than a pop followed by a push.
result = heapq.heapreplace(h, 8)
print(result)Finding Largest and Smallest
The module provides functions to find the n largest or smallest elements in a dataset.
nlargest()
import heapq
nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
print(heapq.nlargest(3, nums))
# Output: [42, 37, 23]nsmallest()
print(heapq.nsmallest(3, nums))
# Output: [-4, 1, 2]These functions also accept a key argument for more complex data structures.
portfolio = [
{'name': 'IBM', 'shares': 100, 'price': 91.1},
{'name': 'AAPL', 'shares': 50, 'price': 543.22},
{'name': 'FB', 'shares': 200, 'price': 21.09},
{'name': 'HPQ', 'shares': 35, 'price': 31.75},
{'name': 'YHOO', 'shares': 45, 'price': 16.35},
{'name': 'ACME', 'shares': 75, 'price': 115.65}
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])