Partial sorting in Python

Why this matters

Sometimes you do not need the whole list in order. A teacher with 500 test marks wants the best 3. A shop wants the 10 cheapest products on a page. A game wants the top 5 scores. Sorting all 500 marks to read the top 3 does more work than the question needs.

This is called partial sorting. You sort only the part of the list that answers the question. The most common cases are "the k smallest items" and "the k largest items", where k is a small number you choose.

Python has two tools for this. The built-in sorted() sorts everything, and then you take the first k items. The heapq module has nsmallest and nlargest, which do the partial work for you. This lesson builds the two methods by hand first, a selection loop and a bounded heap, so you can see how the work is done. It then checks them against a full sort, and shows the mistakes that show up most often.

The earlier lessons in this module covered sorted() and the comparison rules. The heap sort lesson explained the heap idea that the heapq module uses. This lesson uses that idea as a tool.

The idea in plain words

Imagine a teacher with ten marks on paper and a request for the best three. There are two simple ways to do it.

The first way is to scan the whole pile, pick the highest mark, put it aside, and repeat three times. That is a selection loop. Each round looks at the remaining marks once.

The second way is to keep a short list of the best three marks seen so far. For each new mark, compare it with the weakest mark on the short list. If the new mark is better, swap it in. At the end, the short list holds the best three. That is a bounded list, and a heap is a good way to keep it, because it can always tell you which item is the weakest.

Both ways give the same answer. They differ in how much work they do, and in how much code they need.

The selection loop

A selection loop is the simplest partial sort. It runs k rounds. In each round it finds the smallest value that has not been placed yet, and moves it to the front of the list. After k rounds, the first k places hold the k smallest values, in order.

Here is a complete program that uses a selection loop and a bounded heap on the same list. It also tests the edge cases k equal to 0 and k larger than the number of values. The data is built in, so you can run it without typing input.

The output is:

selection, 3 smallest: [1, 3, 7]
heap, 3 smallest:      [1, 3, 7]
full sort, first 3:    [1, 3, 7]
top 3 players: [(91, 'Ravi'), (91, 'Kiran'), (88, 'Meera')]
k = 0 by heap: []
k larger than n by heap: [2, 5]

Start with the selection method. The list is [42, 7, 19, 3, 88, 7, 25, 1, 60, 13], and k is 3. The method copies the list first, so the caller's list is not changed. This matters because a sorting function that changes its input can surprise the caller later.

In the first round, the loop starts with best at index 0, which holds 42. It checks every later value and finds 1 at index 7. The code swaps index 0 and index 7, so the list now starts 1, 7, 19, 3, 88, 7, 25, 42, 60, 13.

In the second round, the loop starts at index 1, which holds 7. It finds 3 at index 3, and swaps it into index 1. The list starts 1, 3, 19, 7, 88, 7, 25, 42, 60, 13.

In the third round, the loop starts at index 2, which holds 19. It finds the first 7 at index 3 and swaps it into index 2. The list starts 1, 3, 7, 19, 88, 7, 25, 42, 60, 13. The loop stops after three rounds, and the function returns the first three values, [1, 3, 7].

The cost is easy to count. Each round looks at every value that has not been placed yet. For n values and k rounds, the loop does about n times k comparisons. When k is small, such as 2 or 3, this is cheap. When k is close to n, it is almost the same as a full selection sort, which is slower than the built-in sort.

The line min(k, len(copy)) handles the case where k is bigger than the list. Without it, the loop would try to place a value at an index past the end of the list. Remember the rule: clamp k to the length before you start.

The bounded heap

A heap is a structure that keeps the smallest item at the top, in a way that makes adding and removing fast. Python's heapq module gives you a min-heap, and the functions heapq.heappush and heapq.heappop add an item and remove the smallest one. The heap is stored inside an ordinary list. You do not need the internal details. You need two facts: heappush and heappop are fast, and the item at index 0 is always the smallest.

heapq has only a min-heap. To get a max-heap, store the negated numbers. The largest number becomes the most negative, so the smallest item in the heap is the largest value you started with. The function smallest_by_heap does this.

A bounded heap keeps at most k items. The rule for the k smallest values is this:

1. Use a max-heap, which here is a min-heap of negated values. Its top item is the largest value kept so far.
2. For each value in the list, push it onto the heap.
3. If the heap now holds more than k items, pop the top item. That is the largest, so it is the one that does not belong among the k smallest.

At the end, the heap holds the k smallest values. The smallest_by_heap function does exactly this.

Tracing the heap by hand

Trace the function on the same list, with k equal to 3. The table shows the heap contents after each step, as plain sets of values. The internal order of a heap is not the same as a sorted order, so the table does not show it.

At the end the heap holds 1, 3 and 7. Those are the three smallest values in the list. Notice that a value can be added and then removed on the same step, such as 88 and 60. A value that is too big for the set of k smallest never stays in the heap.

The heap method does one push and at most one pop for each value. Each push and pop costs about log k steps, because the heap only holds k items. So the whole method takes about n times log k steps. Compare this with a full sort, which takes about n times log n steps. When k is much smaller than n, the bounded heap does less work. When k is close to n, the saving is small, and a full sort is often simpler.

The k largest values

To keep the k largest values, flip the rule. Use a min-heap that holds the values as they are. Its top item is the smallest value kept so far. When a new value is bigger than the top item, the top item does not belong among the k largest, so remove it. The function top_players in the program uses this idea with tuples.

heap = []
for name, score in players:
    heapq.heappush(heap, (score, name))
    if len(heap) > k:
        heapq.heappop(heap)

The heap ends up with the k players who have the highest scores. The function then sorts the heap's items in descending order, so the best player comes first. Two things are worth noting.

First, a heap does not keep its items in sorted order when you print the list or loop over it. The only guaranteed order is the one you get by popping items or by sorting a copy. The function sorts the heap's items to give the output in a clear order.

Second, the program has two players with 91 points. Both are in the top three, so the boundary is not a tie problem here. A tie matters when two players share the k-th score, because then only one of them can stay in the heap. The tuple (score, name) breaks the tie by name. The smaller name is removed first, so the alphabetically later name stays. The result is fixed, and it does not depend on the order of the input.

The heapq module also has nlargest and nsmallest, which do this work in one call. The next program shows the same answers from the library:

The output is:

nsmallest: [1, 3, 7]
nlargest: [88, 60, 42]
nlargest players: [('Ravi', 91), ('Kiran', 91), ('Meera', 88)]

The key argument tells the function which part of each item to compare. Here it is the score, which is the second item of each pair. The library method does the same work as the hand-written heap, and it is the better choice in ordinary code. The hand-written version shows the idea.

Checking the result against a full sort

A partial sort is easy to get subtly wrong, so test it. The simplest test is to compare its answer with the answer of a full sort on many random lists. If both match on every trial, you have good evidence that the method is right. The program below makes 500 random lists with a fixed seed, so the same lists appear on every run. A seed is a starting number for a random generator. The same seed always gives the same sequence of numbers.

The output is:

checks: 500 mismatches: 0

The trials use values from 0 to 49, so many values repeat. They also use k values from 0 up to n + 2, which covers k equal to 0, k inside the list, and k bigger than the list. Repeated values and the edge cases are where partial methods usually fail, so a good test includes them.

The expected answer is found with sorted, then taking the first k values, clamped to n. That is the reference answer. The two partial methods must match it exactly.

The test does not prove the methods are right for every possible input. It shows they agree on 500 varied cases. If a later change breaks one method, the test is likely to show it.

Common mistakes

These are the mistakes that appear most often in partial sorting. The first three come with a program and its real output.

Using the wrong heap direction

To keep the k smallest values, you need a max-heap, so that the largest kept value is removed first. A min-heap removes the smallest value, which keeps the largest values instead. The output is silently wrong, with no error. The program below uses a min-heap for the k smallest values and prints the result. It then prints the internal order of a heap, and the order that heappop gives.

The output is:

min-heap kept: [42, 60, 88]
heap list printed: [88, 60, 25, 42, 13, 7, 19, 1, 3, 7]
pop order: [88, 60, 42, 25, 19, 13, 7, 7, 3, 1]

The first line is the silent error. The method kept 42, 60 and 88, which are the three largest values, not the three smallest. Nothing crashed, so the bug would only show up in a test that checks the answer. This is why the random check above is useful.

The second line shows the heap's list directly. The order 88, 60, 25, 42, ... is the order of the internal list, not the sorted order. Only heappop gives values in order. The pop order line shows the values coming out from largest to smallest, because this heap holds negated numbers. If you need ascending order, collect the values and sort them.

A loop that runs past the end of the list

The selection loop below runs five rounds on a three-value list. In round four, i is 3, and the code reads a[3], which does not exist. Python stops the program with an IndexError:

The output is:

Traceback (most recent call last):
  File "/path/to/loop.py", line 7, in <module>
    a[i], a[best] = a[best], a[i]
                    ~^^^^^^
IndexError: list index out of range

The last line says an index was too large for the list. The fix is to clamp k to the list length first, as the function smallest_by_selection does with min. Python does not treat a negative index as an error, so the loop above fails only because the index is too large. A negative index would count from the end of the list, which is a quieter bug.

Using objects that have no order

A heap of objects needs them to compare with <. The Item class in the next program does not define an order. The first push works, because a heap with one item needs no comparison. The second push compares the new item with the first one, and that fails:

The output is:

Traceback (most recent call last):
  File "/path/to/items.py", line 9, in <module>
    heapq.heappush(items, Item(1))
    ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
TypeError: '<' not supported between instances of 'Item' and 'Item'

The fix is to define __lt__ on the class, or to store tuples such as (weight, count, item) in the heap, where count is a unique number, so the weight decides the order. The Comparison contracts lesson in this module explains how to write __lt__.

Forgetting that tie values are not ordered

If several values are equal at the boundary, the heap keeps only some of them. For plain numbers this does not matter, because equal numbers look the same. For objects, it can matter a lot. Decide the tie rule first, then write it into the key. The players example above shows the idea, with the name as the tie rule.

Changing the input list

The selection function in this lesson copies the list before it starts. A function that sorts in place changes the caller's list, which can break other code that expects the original order. If you do not need the original, sort in place and say so in the function name.

Forgetting the k equal to 0 case

With k equal to 0, the heap method pushes one value and pops it at once, so the result is empty. That is correct. A selection method with k equal to 0 runs no rounds and returns an empty list. Both are fine, but test the case, because an off-by-one error often shows up here first.

Practice

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

1. Find the three largest values in [4, 9, 1, 9, 6]. Use a bounded min-heap of size 3, and return the values in ascending order. Write the function, then predict the output before you run it.

Show the answer

import heapq
 
def largest_three(a):
    h = []
    for x in a:
        heapq.heappush(h, x)
        if len(h) > 3:
            heapq.heappop(h)
    return sorted(h)

The output is [6, 9, 9]. A min-heap of size 3 always removes the smallest value when it has more than three items. That leaves the three largest. Sorting the heap's items gives them in ascending order.

2. Use a selection loop to find the two smallest values in [5, 1, 4, 1]. Write the answer as a list.

Show the answer

The answer is [1, 1]. The first round finds the first 1 at index 1 and swaps it with index 0, so the list becomes [1, 5, 4, 1]. The second round finds the other 1 at index 3 and swaps it with index 1, so the list becomes [1, 1, 4, 5]. The first two values are 1, 1. The two equal values show that a selection loop handles repeats without any extra code.

3. A student wrote this code to keep the three largest values. It has a bug. Find the bug, and say which values it keeps for [4, 9, 1, 9, 6].

h = []
for x in a:
    heapq.heappush(h, -x)
    if len(h) > 3:
        heapq.heappop(h)

Show the answer

The bug is the heap direction. The values are negated, so the heap is a max-heap, and it removes the largest value when it has four items. That leaves the three smallest values, which is the opposite of the goal. For [4, 9, 1, 9, 6] the code keeps 1, 4, 6. The fix is to remove the minus signs, so the min-heap removes the smallest value, as the function in the practice answer above does.

4. Write a function largest_descending(a, k) that returns the k largest values in descending order, using a bounded heap. For [4, 9, 1, 9, 6] and k equal to 3, the answer is [9, 9, 6].

Show the answer

import heapq
 
def largest_descending(a, k):
    h = []
    for x in a:
        heapq.heappush(h, x)
        if len(h) > k:
            heapq.heappop(h)
    return sorted(h, reverse=True)

The min-heap leaves the k largest values, and sorting them in reverse gives the descending order. The function returns [9, 9, 6] for this input.

5. When is a selection loop a better choice than a bounded heap? Answer in two or three sentences.

Show the answer

A selection loop is a better choice when k is very small, such as 1 or 2, and the code must be as short as possible. It needs no import and no negation, and it is easy to read. A bounded heap is better when k is larger, because the selection loop does about n times k comparisons, while the heap does about n times log k steps. For a large list with a moderate k, the heap is the faster choice.

Summary

• Partial sorting finds the k smallest or largest values without sorting the whole list, which saves work when k is small.
• A selection loop runs k rounds and moves the smallest remaining value to the front each time. It is short and simple, and its cost grows with n times k.
• A bounded heap keeps at most k items. Use negated values in a min-heap for the k smallest values, and a plain min-heap for the k largest, and remove the top item whenever the heap is too big.
• The heapq module has nsmallest and nlargest, which do the same work in one call.
• Clamp k to the list length, and test k equal to 0, k larger than the list and repeated values.
• Check a partial method against a full sort on many random lists, with a fixed seed so the test is the same on every run.