What is Python's heapq module?

minerals picture minerals · Nov 14, 2013 · Viewed 63.2k times · Source

I tried "heapq" and arrived at the conclusion that my expectations differ from what I see on the screen. I need somebody to explain how it works and where it can be useful.

From the book Python Module of the Week under paragraph 2.2 Sorting it is written

If you need to maintain a sorted list as you add and remove values, check out heapq. By using the functions in heapq to add or remove items from a list, you can maintain the sort order of the list with low overhead.

Here is what I do and get.

import heapq
heap = []

for i in range(10):
    heap.append(i)

heap
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

heapq.heapify(heap)    
heapq.heappush(heap, 10)    
heap
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

heapq.heappop(heap)
0    
heap
[1, 3, 2, 7, 4, 5, 6, 10, 8, 9] <<< Why the list does not remain sorted?

heapq.heappushpop(heap, 11)
1
heap
[2, 3, 5, 7, 4, 11, 6, 10, 8, 9] <<< Why is 11 put between 4 and 6?

So, as you see the "heap" list is not sorted at all, in fact the more you add and remove the items the more cluttered it becomes. Pushed values take unexplainable positions. What is going on?

Answer

Martijn Pieters picture Martijn Pieters · Nov 14, 2013

The heapq module maintains the heap invariant, which is not the same thing as maintaining the actual list object in sorted order.

Quoting from the heapq documentation:

Heaps are binary trees for which every parent node has a value less than or equal to any of its children. This implementation uses arrays for which heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2] for all k, counting elements from zero. For the sake of comparison, non-existing elements are considered to be infinite. The interesting property of a heap is that its smallest element is always the root, heap[0].

This means that it is very efficient to find the smallest element (just take heap[0]), which is great for a priority queue. After that, the next 2 values will be larger (or equal) than the 1st, and the next 4 after that are going to be larger than their 'parent' node, then the next 8 are larger, etc.

You can read more about the theory behind the datastructure in the Theory section of the documentation. You can also watch this lecture from the MIT OpenCourseWare Introduction to Algorithms course, which explains the algorithm in general terms.

A heap can be turned back into a sorted list very efficiently:

def heapsort(heap):
    return [heapq.heappop(heap) for _ in range(len(heap))]

by just popping the next element from the heap. Using sorted(heap) should be faster still, however, as the TimSort algorithm used by Python’s sort will take advantage of the partial ordering already present in a heap.

You'd use a heap if you are only interested in the smallest value, or the first n smallest values, especially if you are interested in those values on an ongoing basis; adding new items and removing the smallest is very efficient indeed, more so than resorting the list each time you added a value.