Heap sort in Python

Why this matters

Sorting a list with bubble sort or insertion sort takes about n times n steps when the input is in the wrong order. Quick sort is faster on average, but its worst case can also reach about n times n steps when the pivot choices go badly. Heap sort never does that. Its worst case is O(n log n).

The notation O(n log n) (read "big O of n log n") describes how the work grows as the input grows. If you double the number of items, the work grows a little more than double. It does not grow four times, as it would for n times n.

Heap sort also uses only a few extra variables. The sort rearranges the original list and does not copy it into a second list. That is what "in place" means. Merge sort, by contrast, needs a second list of the same size while it works.

Heap sort uses a heap, the same tree-shaped structure that Python's heapq module uses for a priority queue. The heaps module, later in this course, covers heapq. This lesson is different in one way. The heap lives inside the list you want to sort, so no second structure is needed.

The idea in one picture

A heap is a tree where every parent is at least as large as its children. A max-heap is a heap where the largest value is at the top. Here is a max-heap with five values, drawn as a tree:

        10
       /  \
      5    3
     / \
    4   1

The tree is stored in a list in level order, top row first and left to right. The same five values as a list are:

index:  0   1   2   3   4
value: 10   5   3   4   1

Once the tree is a list, you can find any node's relatives with arithmetic. Index 0 is the top. For any index i, the left child is at 2 * i + 1 and the right child is at 2 * i + 2. The parent of i is at (i - 1) // 2, using whole-number division.

Check this with the list above. The children of index 0 are at indexes 1 and 2, which hold 5 and 3. The children of index 1 are at indexes 3 and 4, which hold 4 and 1. Both of index 1's children are smaller than 5, so index 1 is in the right place. The heap rule holds everywhere.

Heap sort works in two phases. The first phase turns the whole list into a max-heap. The second phase repeatedly takes the top of the heap, which is the largest value, and moves it to the end of the list. Each time, the heap gets one item smaller, and the sorted part at the end grows by one item.

The core step: sift down

Both phases use one operation, called sift down. "Sift" means to move something through a sieve, and here the idea is to let a value settle to the right depth in the tree. Sift down takes one value and moves it down the tree until it is not smaller than either of its children.

Here is the procedure for one value at index i:

1. Look at the left and right children of i. Ignore a child whose index is past the end of the heap.
2. Find the largest of the three values: the value at i and its two children.
3. If the value at i is the largest, stop.
4. Otherwise, swap the value at i with the largest child. Set i to the index of that child and repeat.

Take the list [4, 10, 3, 5, 1]. Start at index 0, which holds 4. Its children are 10 at index 1 and 3 at index 2. The largest of 4, 10 and 3 is 10, so swap 4 and 10. The list is now [10, 4, 3, 5, 1].

The value 4 is now at index 1. Its children are 5 at index 3 and 1 at index 4. The largest of 4, 5 and 1 is 5, so swap again. The list is now [10, 5, 3, 4, 1].

The value 4 is now at index 3. Index 3 has no children, because the children would be at indexes 7 and 8, which are past the end. The sift stops. The list is a valid max-heap.

Here is the sift down function. It takes the list, the starting index and the size of the heap. The size matters because part of the list may already hold sorted values that are not part of the heap.

def sift_down(a, i, size):
    while True:
        left = 2 * i + 1
        right = 2 * i + 2
        largest = i
        if left < size and a[left] > a[largest]:
            largest = left
        if right < size and a[right] > a[largest]:
            largest = right
        if largest == i:
            return
        a[i], a[largest] = a[largest], a[i]    # swap with the larger child
        i = largest

The while True loop repeats until the return runs. The condition left < size is the check from step 1. It stops the code from reading an index that is past the end of the heap. Without that check, Python raises an error, as the mistakes section shows.

The swap on the line a[i], a[largest] = a[largest], a[i] is the one-line swap from the first lesson of this module. Python needs no temporary variable for it.

Each swap moves a value one level down the tree. A tree with n items has about log n levels, so a sift down does at most about log n swaps. This is where the log n in the total cost comes from.

Step 1: build the max-heap

The list [4, 10, 3, 5, 1] is not a heap yet. For example, the root is 4, but index 1 holds 10, which is larger than its parent.

To build the heap, call sift down on every index that has at least one child, working from the last such index back to the root. An index i has a child when 2 * i + 1 is less than the size. For a list of length n, the last index with a child is n // 2 - 1.

Why start from the back? Every index after n // 2 - 1 is a leaf, which means it has no children. A leaf is already a heap of one value. When you reach index i, both of its children's subtrees are already heaps, so one sift down is enough to fix the subtree rooted at i.

Here is a function that builds the heap and prints the list after each sift, so you can see the work:

No input is needed. The output is:

before: [4, 10, 3, 5, 1]
  after sift at index 1: [4, 10, 3, 5, 1]
  after sift at index 0: [10, 5, 3, 4, 1]

The first line of work, at index 1, changes nothing. The value 10 is already larger than its children 5 and 1. The second line, at index 0, does the real work, as traced in the previous section. After the loop, the list is a max-heap.

The range range(len(a) // 2 - 1, -1, -1) counts down from the last index with a child to 0. The third number in range is the step, and -1 means "one step backwards." The end value -1 is not included, so the loop stops at 0.

Step 2: move the top to the end

Now the largest value, 10, sits at index 0. The sort phase repeats three actions. Each round:

1. Swap index 0 with index end, the last position of the heap. The largest value goes to its final place.
2. Shrink the heap by one. The heap now covers indexes 0 to end - 1.
3. Sift down index 0 inside the smaller heap, so the list is a max-heap again.

The variable end starts at the last index and moves left by one each round. The loop stops when end is 0, because a single value left in the heap is already in its place.

Here is the whole sort, with the list printed after each swap and sift. It uses the same sift_down and swap functions from the last example.

No input is needed. The output is:

before: [4, 10, 3, 5, 1]
  top moved to index 4: [5, 4, 3, 1, 10]
  top moved to index 3: [4, 1, 3, 5, 10]
  top moved to index 2: [3, 1, 4, 5, 10]
  top moved to index 1: [1, 3, 4, 5, 10]
after:  [1, 3, 4, 5, 10]

Read the first line of the sort phase closely. Before the swap, the list was [10, 5, 3, 4, 1]. Swapping index 0 with index 4 gives [1, 5, 3, 4, 10]. The 10 is now at its final place. The heap now covers indexes 0 to 3, which hold 1, 5, 3, 4. Sift down moves 1 past 5, and then past 4, leaving [5, 4, 3, 1] in the heap part. The list shows [5, 4, 3, 1, 10].

Each later line works the same way. The sorted part at the end grows from 10 to 5, 10, then 4, 5, 10, and so on, until the whole list is sorted.

Heap sort leaves its answer in the original list. The after line prints that same list.

A second example: any type, with a key

The first version works only for numbers, because the rule "larger is better" is written directly into the code with >. A more useful version takes a key function, which gives the value to compare for each item. The heap then works on any objects that the key can compare.

The next program makes the heap sort take a key. It checks the sort on 2000 random lists against Python's sorted(). It also sorts a list of players by score. The player class is a dataclass, which is the short way to write a class that only holds fields.

No input is needed. The output is:

random trials failing: 0
ascending by score: [Player('Ben', 70), Player('Dee', 85), Player('Cy', 90), Player('Eli', 90), Player('Ana', 90)]
empty length after sort: 0
words: ['fig', 'pear', 'apple']

The first line is the test result. Zero failures in 2000 random lists, including lists of length 0 and 1, and lists with repeated and negative values, is good evidence that the sort is correct. It is not a proof, but it catches most slips in the indexes.

The function heap_sort(a, key=...) has a default key, lambda x: x, which returns each item unchanged. When you call heap_sort(mine) with no key, the items are compared as they are. When you pass key=lambda p: p.score, the heap compares scores. The key is called inside the sift, so it is called many times. Keep it cheap.

The last word example uses a key that returns a tuple. Python compares tuples one item at a time, so (len(w), w) sorts by length first and then alphabetically. The result is fig, pear, apple, which puts the shortest word first.

A third example: heap sort is not stable

A sort is stable when items that compare as equal keep their original order. Stability matters when you sort a list twice, first by one field and then by another, or when a list of records has a second field that should keep its order.

Look at the players output again. The input had Ana, Cy and Eli, all with score 90. Their order in the input was Ana, then Cy, then Eli. After the heap sort, the order of those three is:

Cy, Eli, Ana

The scores are sorted correctly, but the three 90s no longer keep their original order. Heap sort is not stable. The swaps move items past one another, even when they are equal.

If you need a stable sort, use Python's sorted() or list.sort(). They use a stable merge-based method, so equal items keep their order. The next lessons in this module cover multi-key sorting, which depends on stability. Heap sort is the right choice when you need a guaranteed O(n log n) time and little extra memory, and you do not need equal items to keep their order.

Python's heapq module

Python's heapq module gives you a heap that is ready to use. It is a min-heap, which means the smallest value is at the top. The functions heapq.heappush and heapq.heappop add and remove values, and heapq.heapify turns a list into a heap in place.

The output is:

heap: [1, 3, 2]
smallest first: [1, 2, 3]

heapq is a min-heap, so to get the largest value first, store the negated numbers and negate them again when you read them. The example in the practice section uses this trick. The heapq module keeps its own list, so using it is not in-place heap sort. It takes extra memory for a second list, which is the cost the common mistakes section describes.

Common mistakes

Forgetting the size check when sifting

The most common bug is reading the right child without checking that it exists. For the list [1, 2, 3, 4], the right child of index 1 is at index 4, which is past the end. Here is a version of sift down with the check removed from the right child:

Calling it on that list raises this error:

Traceback (most recent call last):
  File "heap.py", line 12, in <module>
    sift_down_no_check([1, 2, 3, 4], 1, 4)
    ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
  File "heap.py", line 7, in sift_down_no_check
    if a[right] > a[largest]:
       ~^^^^^^^
IndexError: list index out of range

Read the last line first. IndexError: list index out of range means the code asked for an index that is not in the list. The line above it names the code that asked. The file name and line numbers depend on where you saved the program. The fix is the same as in the sift down function above: test right < size before reading a[right].

Notice the difference from negative indexes. A negative index in Python does not raise an error, because it counts from the end of the list. An index that is too large does raise an error. Here the index 4 is too large, so the error appears.

Shrinking the heap in the wrong place

Heap sort works only if the sift in the sort phase uses the shrinking end as the size. If it uses the full length of the list, the sorted values at the end come back into the heap. Here is a version that passes len(a) in the second loop:

Running it on [4, 10, 3, 5, 1] does not raise an error, and the output is wrong:

wrong size result: [10, 5, 3, 4, 1]

The list is not sorted. The sift keeps moving values that are already in their final places. Nothing crashes, so this kind of bug is easy to miss. A quick check with a few small lists, comparing the result to sorted(), catches it.

Expecting heap sort to be stable

The players example shows that equal items can change order. If your program depends on the original order of equal items, heap sort is the wrong tool. Use sorted() instead, as described above.

Calling heapq "in place"

A common shortcut is to push every value into a heapq heap and then pop until the heap is empty. That works, and the output is sorted. It is not the in-place heap sort from this lesson, because the heapq heap keeps its own list. That list needs memory for n values. Extra memory of O(n) size is more than the O(1) that this lesson's version uses.

Getting the order backward

The code in this lesson sorts smallest to largest. A max-heap naturally puts the largest value at the top, so each swap leaves the largest at the end, and the list ends up in ascending order. To sort from largest to smallest, you need a min-heap, where the smallest value is at the top. The practice section gives the change.

Practice

Try each exercise in a file first. Check your answer after you have made an attempt.

Exercise 1: build a heap by hand

Start with the list [1, 3, 2]. Build a max-heap using only the first step of heap sort. Write the list after the build.

Show the answer

The list has three values, so the last index with a child is 3 // 2 - 1, which is 0. Only index 0 needs a sift. The children of index 0 are 3 at index 1 and 2 at index 2. The largest of 1, 3 and 2 is 3. Swap 1 and 3. The list is [3, 1, 2]. Index 1 has no children, so the sift stops.

Exercise 2: why the loop stops at 1

In the sort phase, the loop is for end in range(len(a) - 1, 0, -1). Why does it stop before end reaches 0?

Show the answer

After the round with end equal to 1, the heap holds only index 0, and that single value is already in its final place. The round with end equal to 0 would swap index 0 with itself and sift an empty heap, so it changes nothing. The loop skips that useless round.

Exercise 3: sort from largest to smallest

Change the heap sort so that it sorts from largest to smallest. The result for [4, 10, 3, 5, 1] should be [10, 5, 4, 3, 1].

Show the answer

Change every > in the sift down function to <. The function then moves the smallest child up, so the top of the heap is the smallest value. Moving that top to the end leaves the smallest value at the end, and the list ends up in descending order. Here is the sift function after the change:

def sift_down(a, i, size):
    while True:
        left = 2 * i + 1
        right = 2 * i + 2
        smallest = i
        if left < size and a[left] < a[smallest]:
            smallest = left
        if right < size and a[right] < a[smallest]:
            smallest = right
        if smallest == i:
            return
        a[i], a[smallest] = a[smallest], a[i]
        i = smallest

The build loop and the sort loop stay the same. This version was run on the list and printed [10, 5, 4, 3, 1].

Exercise 4: empty and one-value lists

What does heap sort do with a list of length 0 and a list of length 1? Does it raise an error?

Show the answer

It does nothing and raises no error. For length 0, the build loop starts at index 0 // 2 - 1, which is -1, so the loop does not run. The sort loop starts at -1 too, so it does not run either. For length 1, the build loop starts at -1 again, and the sort loop starts at end = 0, so range(0, 0, -1) gives no values. Both lists are returned unchanged.

Exercise 5: find the bug

A student wrote this sort phase and says it sorts correctly because the build is correct:

for end in range(len(a) - 1, 0, -1):
    swap(a, 0, end)
    sift_down(a, 0, len(a))

Run it on [4, 10, 3, 5, 1]. The result is [10, 5, 3, 4, 1]. Explain why.

Show the answer

The sift passes len(a) as the size, so the heap always covers the whole list. The value swapped to the end is part of the heap again, and the sift can move it back to the front. The sorted part at the end is not protected. The fix is to pass end as the size, so the sift stops before the sorted part.

Exercise 6: a max-heap with heapq

Use heapq to print the largest value of [4, 10, 3, 5, 1] first, then the rest in descending order.

Show the answer

Push the negated values, then pop them. Each popped value is negated back.

import heapq
 
h = []
for v in [4, 10, 3, 5, 1]:
    heapq.heappush(h, -v)
print([-heapq.heappop(h) for _ in range(5)])

The output is [10, 5, 4, 3, 1]. The smallest negated value is the largest original value, so the min-heap gives the largest first.

Summary

• Heap sort turns the list into a max-heap, then moves the top to the end, one position at a time.
• A heap is stored in a list. The children of index i are at 2 * i + 1 and 2 * i + 2.
• The sift down step moves one value down the tree and costs about log n swaps.
• The build phase runs from index n // 2 - 1 back to 0. The sort phase shrinks the heap by one each round.
• The total time is O(n log n) in the best and worst cases, and it uses O(1) extra space.
• Heap sort is not stable, so equal items can change their order.