Top-k with heapq

Why it matters

Many programs need only a few of the biggest items from a large collection. A shop wants the five best-selling products from a million sales records. A game wants the ten highest scores from all players. A search engine wants the twenty most relevant pages out of millions. In each case the full list is far too long to read, but the answer is short.

The obvious way is to sort the whole list from biggest to smallest and take the first k items. That works, but it does more work than the question asks for. Sorting a million numbers to find the top ten puts nearly a million numbers in order when only ten matter.

This lesson shows a faster way. It keeps a small basket of k items, and the basket is organised so that the weakest item in the basket is always easy to find. The organiser is a heap. A heap is a tree-shaped way of storing items so that the smallest (or largest) one is always at the top. In Python the heap functions live in the heapq module, and they work on an ordinary list. The heaps lesson of this module explains heapq from the start. This lesson takes its basic use for granted and focuses on the top-k method, why it is correct and how it behaves.

The problem: the k largest values

Take the list 5, 1, 9, 3, 7, 9, 2. The top 3 values are 9, 9, 7. Written in order from smallest to largest, the answer is [7, 9, 9]. Note that 9 appears twice, because the list has two 9s. The answer is a multiset, not a set: a value that appears twice counts twice.

The sorting answer and its cost

The sorting answer is short.

1. Sort the list from smallest to largest, giving 1, 2, 3, 5, 7, 9, 9.
2. Take the last 3 items, giving 7, 9, 9.

In Python the whole answer is sorted(values)[-3:]. The cost of a good sort is about n * log n steps, where n is the list length. Here log n means the number of times you can cut n in half before you reach 1. For a list of a million items, log n is about 20. So sorting takes about twenty steps for each item. This is fine for small lists. For a stream of numbers that never ends, sorting is not possible at all, because you cannot sort what you have not seen yet.

A stream is a sequence of values that arrive one at a time, such as sales records coming from a store all day. You want the top k so far, and you do not want to keep every value in memory.

The basket idea

Imagine a basket that holds exactly k items. Start by filling it with the first k values you see. For each later value, compare it with the weakest item in the basket.

• If the new value is bigger than the weakest item, throw the weakest item out and put the new value in.
• If the new value is not bigger, throw the new value away.

At the end, the basket holds the k largest values seen. The method needs only the weakest item of the basket at each step. That is where the heap comes in. A min-heap keeps the smallest item at the top, so the weakest item can be read in one step.

The idea: a size-k min-heap

A min-heap is a heap where the smallest item is at the top. In Python, heapq keeps a min-heap inside a list, and the smallest item is always at index 0. You read the top with heap[0], remove it with heapq.heappop(heap), and add a new item with heapq.heappush(heap, x). The function heapq.heapreplace(heap, x) removes the top and adds x in one call, which is what the top-k method needs.

The output is 1, because 1 is the smallest item. Adding and removing take about log k steps each, where k is the number of items in the heap. The heap never grows past k items if the code is written as shown below.

The three rules

Each value in the stream follows one of three cases.

1. If the heap has fewer than k items, add the value.
2. If the heap is full and the value is bigger than the top (the smallest item), remove the top and add the value.
3. Otherwise, do nothing. The value is too small to be in the top k.

Here is the method.

import heapq
 
 
def top_k(values, k):
    if k == 0:
        return []
    heap = []
    for v in values:
        if len(heap) < k:
            heapq.heappush(heap, v)
        elif v > heap[0]:
            heapq.heapreplace(heap, v)
    return sorted(heap)

The if k == 0 check at the top is needed. The reason is explained in the mistakes section below. The last line sorts the result, because the heap only keeps the top at the front. The items inside the heap are not in order.

A worked example, one value at a time

Run the method on 5, 1, 9, 3, 7, 9, 2 with k = 3. The table shows the heap after each value. The table lists the heap in sorted order, so the smallest item is first. The heap's own list may have a different order inside, but its top is always the smallest item.

The heap ends as [7, 9, 9]. The output of the method is the same list, sorted. The value 2 was skipped at the end, and the value 1 was thrown out in the fourth row. Both are correct, because neither is in the top three.

Running the code gives:

The output is:

top 3 of [5, 1, 9, 3, 7, 9, 2]: [7, 9, 9]
top 0: []
top 10 (k larger than n): [1, 2, 3, 5, 7, 9, 9]

The third line shows what happens when k is bigger than the list. The heap never fills, so every value is added, and the result is the whole list sorted.

Why the cost is n log k

The method reads each of the n values once. For each value, it does at most one replace on a heap of at most k items. Each of those costs about log k steps. The total is about n * log k steps.

Compare the two costs for a list of a million numbers with k = 10.

• Sorting: about 1,000,000 * 20 steps.
• Top-k with a heap: about 1,000,000 * 3 steps, because log2 10 is about 3.3.

The heap method is faster when k is small. When k is close to n, the two costs get close, and sorting is simpler. Use the heap when k is much smaller than the list length. The standard library also has heapq.nlargest(k, values), which does this work for you. The loop above shows the idea behind it.

The memory is also small. The heap holds k numbers, not the whole list. For a stream of a million numbers and k = 5, the heap holds 5 values.

Worked example two: the k-th largest in a stream

The second example asks a different question. Instead of all top k values, it asks for one value: the k-th largest. The k-th largest is the value that has exactly k - 1 values above it. For k = 1 it is the biggest value. For k = 2 it is the second biggest, and so on.

The same basket works. After the loop, the top of the min-heap is the weakest item in the basket. That item is exactly the k-th largest value in the whole stream, because the basket holds the k biggest values and the smallest of those is the k-th biggest.

The output is:

stream 3rd largest: 16

Take the stream 4, 8, 15, 16, 23, 42, 8, 4, 15. The biggest three values are 42, 23, 16. The third largest is 16. The function returns heap[0] instead of the whole sorted list, so it does not need the sort step.

A detail matters here. The function returns the top of the heap after all values are read. It must not return early. The heap is only correct for the k-th largest when every value has been seen. If the list is empty, heap[0] raises IndexError. If it has fewer than k values, heap[0] returns the smallest value seen, which is not the answer, so check the length first in real code.

Worked example three: the k closest points

The third example is harder, and it reverses the heap. So far we kept the biggest items, using a min-heap. Sometimes we want the smallest items. Take the points on a map, each a pair of x and y numbers. The question is: which k points are closest to the origin, the point (0, 0)?

The distance of a point is x * x + y * y, squared. Squaring keeps the comparisons correct without a square root, because a bigger squared distance always means a bigger distance.

To keep the k closest points, the basket must throw out the farthest point when a closer one arrives. So the heap's top must be the farthest point in the basket. That is a max-heap, a heap with the biggest item on top. heapq has no max-heap, so the trick is to store the negative of the distance. The smallest negative is the largest distance, so the top of the min-heap is the farthest point.

The output is:

2 closest points: [(1, 1), (0, 2)]

Take the points (3, 3), (5, -1), (-2, 4), (1, 1) and (0, 2). Their squared distances are 18, 26, 20, 2 and 4. The two closest are (1, 1) and (0, 2).

Each heap entry is a pair (-d, p). The heap compares entries by their first item, so it orders them by the negated distance. The function -heap[0][0] turns the top's negated distance back into a real distance, and the new point replaces the top only when it is closer.

Notice the two directions. The top-k for biggest values used a min-heap, and the top-k for closest points uses a min-heap of negated distances. The rule is the same in both cases: keep the basket of the k best items, and compare each new item with the weakest one in the basket. "Weakest" means the smallest for the biggest-values problem and the farthest for the closest-points problem.

Checking the method against a slow answer

A good habit is to test a fast method against a slow one that is easy to trust. The scratch test for this lesson generated 5,000 random lists, with random values between -20 and 20, and random k values. For each list it compared the heap method with the sorting answer, once for the top-k list and once for the k-th largest value. The test found no mismatches. The sorting answer is simple enough to be correct by inspection, so it makes a good reference.

The output of that test was:

random mismatches: 0

Common mistakes

Using a max-heap for the biggest values. A beginner often stores the values with the wrong sign, so the heap keeps the wrong end. Here the method stores negated values and then compares a new value with the wrong end of the basket. The test v > -heap[0] asks whether v beats the biggest item kept, so almost no value gets in:

The output is:

wrong: [1, 5, 9]

The answer is [1, 5, 9], the first three values sorted. It is not the three biggest, which are 7, 9, 9. Read the top of the heap as "the item to throw out first". For the biggest values, that must be the smallest item, so the heap must hold the values as they are, not negated.

Forgetting the k = 0 case. When k = 0, the heap is never allowed to hold anything, but the first value still reaches the elif line. On that empty heap, heap[0] has no item to read, and Python raises an error. The next program shows the error without the guard:

The output is:

caught: list index out of range

The message names the problem: the code asked for an item at index 0 of an empty list. The fix is an early return for k == 0, as in top_k.

Printing the heap and assuming it is sorted. The heap keeps only the top in the right place. The other items are in a tree order that looks random when printed. Here is the heap printed directly with k = 5 on the same list:

The output is:

printed heap: [3, 5, 9, 9, 7]
sorted: [3, 5, 7, 9, 9]

The top, 3, is the smallest, and the rest are not in order. Always sort a copy before you show the result to a reader. The method above does that with sorted(heap).

Comparing values with is. The == operator compares values, and the is operator compares identity, which means whether two names point to the same object. Python keeps one shared object for small integers, so is can look right for small numbers. For larger numbers it fails:

The output is:

True False

Both names hold the value 1000, but is says they are different objects. Use == to compare numbers. Use is only for None and for identity checks.

Changing an item that is inside the heap. The heap only checks the order when you push or pop. If the code changes a value that is already inside the heap, the heap is no longer correct. Items in the heap should be treated as read-only. If a value must change, remove it, change it, and add it again. The heaps lesson of this module shows the lazy deletion method for this case.

Using heappop where heap[0] is meant. In the kth_largest function, the answer is heap[0], which reads the top and leaves it in the heap. A beginner might call heapq.heappop(heap) instead. That removes the top from the heap, so the heap changes and any later reading gives a different item. Use heap[0] when you only want to read the answer.

Practice

Exercise 1. Use top_k by hand on [4, 4, 4, 1] with k = 2. What is the output?

Show the answer: the heap takes 4 and 4, which fills it. The next 4 is not bigger than the top 4, so it is skipped. The 1 is smaller than the top, so it is skipped. The output is [4, 4].

Exercise 2. Change top_k so that it returns the k smallest values. Test it on [5, 1, 9, 3] with k = 2.

Show the answer: use a heap of negated values, so the top is the biggest of the basket. A new value replaces the top only when it is smaller:

import heapq
 
 
def k_smallest(values, k):
    if k == 0:
        return []
    heap = []
    for v in values:
        if len(heap) < k:
            heapq.heappush(heap, -v)
        elif v < -heap[0]:
            heapq.heapreplace(heap, -v)
    return sorted(-x for x in heap)

On [5, 1, 9, 3] with k = 2, the basket goes [5], then [5, 1]. The value 9 is not smaller than the top 5, so it is skipped. The value 3 is smaller than 5, so 5 is removed and 3 is added. The output is [1, 3].

Exercise 3. The check v > heap[0] uses >. Would >= also give a correct answer? Explain what changes.

Show the answer: yes, the list of values is the same either way. When v equals the top, replacing the top with v puts back an equal value, so the basket does not change in a way that matters. The only difference is one extra replace for each tie. Use > because it does less work and is easier to read.

Exercise 4. A stream has one million numbers, and you need the top 5. About how many numbers does the heap hold at its largest, and what does that mean for memory?

Show the answer: the heap holds at most 5 numbers at any time. The rest of the million values are read once and then discarded. Memory use does not depend on the stream length. This is why the method works for streams that never end.

Exercise 5. Find the second largest value in [2, 7, 1, 8, 2, 8] using kth_largest with k = 2.

Show the answer: the biggest value is 8, and it appears twice. Counting duplicates, the second largest is also 8. The answer is 8. A value that appears twice takes two places in the ranking.

Summary

• Top-k means keeping the k biggest values seen so far in a basket of at most k items.
• A size-k min-heap keeps the smallest basket item on top, so the weakest item is cheap to read and replace.
• The method costs about n * log k steps, which beats sorting when k is much smaller than n.
• For the k-th largest value, read heap[0] after the loop.
• For the k closest or smallest values, use negated values, so the farthest or biggest item is on top.
• Sort a copy before you show the result, because the heap does not keep its items in order.