Top-k selection in Python: stop sorting values you will discard

Problem. A report often needs only the largest k scores from n values, but the obvious sorted(values, reverse=True)[:k] orders every value before throwing most of them away.^[2] Baseline complexity. The sort has O(n log n) time in current CPython, followed by a slice that copies k references; it also holds the complete sorted list.^[4][2] Solution. heapq.nlargest(k, values) keeps a heap of the current candidates, scans the input, and sorts only the retained candidates at the end.^[1][5] Measured result. On 200,000 seeded integers, the heap version took 9.249 ms for k=20 against 126.289 ms for the sort, a 13.65x median difference; at k=50,000 it took 587.100 ms against 128.807 ms, so the heap lost when k became large.

Reproduce the result

Complete code

from __future__ import annotations
import heapq
import platform
import random
import statistics
import sys
import timeit
import tracemalloc
def top_k_sorted(values: list[int], k: int) -> list[int]:
    """Return the k largest values by sorting the complete input."""
    return sorted(values, reverse=True)[:k]
def top_k_heap(values: list[int], k: int) -> list[int]:
    """Return the k largest values while keeping only a k-item heap."""
    return heapq.nlargest(k, values)
def benchmark(function, values: list[int], k: int) -> list[float]:
    timer = timeit.Timer(lambda: function(values, k))
    return timer.repeat(repeat=7, number=3)
def peak_traced_kib(function, values: list[int], k: int) -> float:
    tracemalloc.start()
    function(values, k)
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return round(peak / 1024, 1)
def main() -> None:
    count = 200_000
    seed = 20260909
    values = random.Random(seed).randrange
    data = [values(0, 10_000_000) for _ in range(count)]
    checks = {}
    edge_cases = {
        "empty": top_k_sorted([], 3) == top_k_heap([], 3) == [],
        "zero": top_k_sorted([4, 1, 9, 2], 0) == top_k_heap([4, 1, 9, 2], 0) == [],
        "larger_than_input": top_k_sorted([4, 1, 9, 2], 10) == top_k_heap([4, 1, 9, 2], 10) == [9, 4, 2, 1],
    }
    timings = {}
    memory = {}
    for k in (20, 5_000, 50_000):
        sorted_result = top_k_sorted(data, k)
        heap_result = top_k_heap(data, k)
        if sorted_result != heap_result:
            raise AssertionError(f"result mismatch for k={k}")
        checks[k] = {
            "equal": True,
            "first": heap_result[0],
            "last": heap_result[-1],
            "length": len(heap_result),
        }
        sorted_samples = benchmark(top_k_sorted, data, k)
        heap_samples = benchmark(top_k_heap, data, k)
        timings[k] = {
            "sorted_ms": [round(sample * 1000 / 3, 3) for sample in sorted_samples],
            "heap_ms": [round(sample * 1000 / 3, 3) for sample in heap_samples],
            "sorted_median_ms": round(statistics.median(sorted_samples) * 1000 / 3, 3),
            "heap_median_ms": round(statistics.median(heap_samples) * 1000 / 3, 3),
            "median_ratio": round(statistics.median(sorted_samples) / statistics.median(heap_samples), 2),
        }
        memory[k] = {
            "sorted_peak_traced_kib": peak_traced_kib(top_k_sorted, data, k),
            "heap_peak_traced_kib": peak_traced_kib(top_k_heap, data, k),
        }
    print(f"python={sys.version.split()[0]}")
    print(f"platform={platform.platform()}")
    print(f"input_count={count}")
    print(f"seed={seed}")
    print("edge_cases=" + repr(edge_cases))
    print("checks=" + repr(checks))
    print("timings_ms_per_call=" + repr(timings))
    print("peak_traced_memory_kib=" + repr(memory))
    print("timing_policy=7 repeats, 3 calls per repeat, median reported")
if __name__ == "__main__":
    main()

Output

python=3.13.5
platform=Linux-6.12.47+rpt-rpi-v8-aarch64-with-glibc2.41
input_count=200000
seed=20260909
edge_cases={'empty': True, 'zero': True, 'larger_than_input': True}
checks={20: {'equal': True, 'first': 9999955, 'last': 9998894, 'length': 20}, 5000: {'equal': True, 'first': 9999955, 'last': 9750777, 'length': 5000}, 50000: {'equal': True, 'first': 9999955, 'last': 7495850, 'length': 50000}}
timings_ms_per_call={20: {'sorted_ms': [147.396, 123.254, 126.289, 127.575, 129.594, 124.692, 122.802], 'heap_ms': [9.506, 9.196, 9.314, 9.176, 9.249, 9.339, 9.2], 'sorted_median_ms': 126.289, 'heap_median_ms': 9.249, 'median_ratio': 13.65}, 5000: {'sorted_ms': [124.529, 123.867, 124.366, 124.16, 124.78, 124.054, 124.069], 'heap_ms': [62.512, 64.649, 82.526, 83.378, 77.285, 80.784, 80.443], 'sorted_median_ms': 124.16, 'heap_median_ms': 80.443, 'median_ratio': 1.54}, 50000: {'sorted_ms': [130.315, 132.96, 128.275, 128.41, 127.803, 128.807, 130.256], 'heap_ms': [554.377, 587.1, 676.509, 604.695, 551.805, 553.549, 650.994], 'sorted_median_ms': 128.807, 'heap_median_ms': 587.1, 'median_ratio': 0.22}}
peak_traced_memory_kib={20: {'sorted_peak_traced_kib': 2343.7, 'heap_peak_traced_kib': 1.2}, 5000: {'sorted_peak_traced_kib': 2343.7, 'heap_peak_traced_kib': 402.3}, 50000: {'sorted_peak_traced_kib': 2343.7, 'heap_peak_traced_kib': 5055.5}}
timing_policy=7 repeats, 3 calls per repeat, median reported

Environment

The run used CPython 3.13.5 on Linux 6.12.47+rpt-rpi-v8, aarch64, with glibc 2.41. The programme uses only the Python standard library.

Methodology

The input contains 200,000 integers from random.Random(20260909).randrange(0, 10_000_000). Each function was checked against the other for k values of 20, 5,000 and 50,000, plus empty input, zero, and k larger than the input. The timing uses timeit.repeat with seven repeats and three calls per repeat; the reported figure is the median per call. Python’s timing documentation describes repeat as repeated calls to the timing loop and notes that its default timer is perf_counter().^[3] Peak memory is traced with tracemalloc for one call, so those numbers cover Python allocations observed by the tracer, not the process’s complete resident memory. The input construction happens outside the timed call.

The problem

Suppose a service stores scores for 200,000 items but the page shows the best 20. A full sort produces a ranking for all 200,000 items. The page then reads the first 20 and leaves the other 199,980 in a list that has already cost time and memory.

That work is sensible when the caller needs the entire ranking. It is a poor match when the output is deliberately small. The useful question is not “which value comes next in the full order?” It is “which values still belong in the best k?”

The two functions in the programme return the same descending list. Their agreement is checked before the stopwatch runs, including cases where the requested count is zero or larger than the input. That matters because a fast wrong answer is still a fairly expensive bug.

Baseline complexity

The baseline is short:

sorted(values, reverse=True)[:k]

sorted() returns a new sorted list.^[2] For a list of n values, the documented CPython complexity reference gives sorting as O(n log n).^[4] The slice adds O(k) work to copy the selected references, but it does not remove the cost of producing the complete order first.

The practical memory picture is similar. The input list remains in memory, and the sorted result temporarily holds n references. The slice creates another list of k references while the expression is being evaluated. The integers themselves are not copied by the sort, but the list storage and sorting workspace still exist.

The solution

A heap stores the smallest candidate at its root when it is a min-heap. Python’s heapq documentation defines that heap invariant and identifies the first heap item as the smallest item.^[1]

For top-k selection, that smallest retained candidate is exactly the one worth replacing. The algorithm starts with k values. For each later value, it compares the value with the heap root. If the new value is larger, the root leaves and the new value enters. If it is smaller, the heap stays put. At the end, only k candidates remain, and those candidates are sorted into the returned descending list.

The CPython 3.13 implementation exposes the important boundary cases: k=1 uses max(), k at least as large as the input uses sorted(), and the general path maintains a heap before sorting the result.^[5] The documentation makes the same practical recommendation: nlargest() is intended for smaller values of n, while sorted() is more efficient for larger ones.^[1]

For 0 < k < n, the general path has an upper-bound cost of O(k) to initialise the heap, O(n log k) for possible replacements, and O(k log k) to order the final candidates. That is commonly written as O(n log k). The exact number of replacements depends on the data, and CPython’s implementation also has shortcuts, so the formula describes the algorithmic shape rather than a promise about every input.

The extra state is O(k), excluding the input and returned list. That is the useful difference for a small result. It is also why the advantage can disappear: as k approaches n, the heap becomes large, its maintenance costs more, and a full sort has a well-optimised path for producing the answer.

What the measurement shows

The three measured sizes tell a cleaner story than one headline speed-up. With k=20, nlargest() was 13.65 times faster by median and traced 1.2 KiB of peak temporary Python allocation, compared with 2,343.7 KiB for the sorting function. With k=5,000, the ratio fell to 1.54 and the traced peaks were 402.3 KiB for the heap path and 2,343.7 KiB for the sort path. At k=50,000, the heap path was 4.56 times slower, while its traced peak reached 5,055.5 KiB.

Those memory figures are not a measurement of total process memory. They describe allocations visible to tracemalloc during one function call, while the input list is already present. The timing figures have their own limits: they come from one machine, one Python build, integer values, one deterministic distribution and a callable passed to timeit. The relative result is useful for this comparison, not a universal ranking of all top-k implementations.

The correctness checks returned equal results for all three measured k values. The edge checks also passed. That verifies the output for this programme’s cases; it does not prove that an arbitrary comparable object, custom key function or streaming source will have the same cost.

Where it stops helping

Use a full sort when the caller needs every item ordered, when k is close to n, or when the measured workload says it is faster. The Python documentation explicitly points readers towards sorted() for larger requested counts.^[1]

For k=1, use max() unless you specifically need the list-shaped interface. CPython already takes that shortcut inside nlargest().^[5]

The simple example also assumes values can be compared directly. Real records usually need a key function, such as heapq.nlargest(k, rows, key=lambda row: row.score). That is still the same selection idea, but key calculation and object comparisons can dominate the heap operations.

Finally, nlargest() returns a list. It does not give a lazy top-k stream, and the input still has to be scanned. If the data arrives continuously, a long-lived fixed-size heap may be a better interface than rebuilding the selection for each request. If the source is already sorted or stored in a database, pushing the operation closer to that source may matter more than choosing between these two Python expressions.

For a small requested result, the practical rule is plain: measure nlargest() against a full sort at the k values your application actually serves. The winning line in this run was the heap for 20 results, and the winning line was the sort for 50,000.

Sources

[1] heapq — Heap queue algorithm — Python 3.13 documentation

[2] Built-in Functions — Python 3.13 documentation

[3] timeit — Measure execution time of small code snippets — Python 3.13 documentation

[4] TimeComplexity — Python Wiki archive

[5] CPython 3.13 Lib/heapq.py