Partial Sort, in depth: bounded heaps, select-then-sort, partial quicksort, library calls and sharded top-N

The problem, precisely

Given n items and a comparison, a partial sort with parameter k rearranges the input so that positions 0 to k-1 hold the k smallest items in ascending order. Every later position holds an item at least as large as position k-1, in unspecified order. Three nearby problems are easy to confuse with it, and mixing them up leads to the wrong tool:

• Selection (the k-th smallest, or the k smallest in any order) does not sort the winners. Quickselect and C++ std::nth_element solve this, with linear expected time.
• Partial sort also orders the winners. It costs selection plus O(k log k) at most.
• Heavy hitters (the k most frequent keys) is a counting problem, not an ordering problem. Read top-K and heavy hitters for sketches that solve it.

Counting arguments set the lower bound. Any comparison algorithm must look at each item at least once, which takes n - 1 comparisons. Ordering k winners takes about log2(k!) comparisons, roughly k log2 k. So Omega(n + k log k) is the target. Heap selection gets close on random input, and selection followed by sorting the prefix meets it in expectation on any input order.

Four strategies and their costs

There are four reasonable strategies. The table gives their costs with k much smaller than n, which is the usual case.

The heap method has a useful property on random input. Element i becomes one of the k smallest seen so far with probability k/i. So the expected number of heap replacements is about k ln(n/k), which is small. Most elements cost one comparison against the root and are dropped. Sorted or reverse-sorted input removes that advantage, as the worked example shows.

The bounded heap in code

To keep the k smallest items, use a max-heap. Its root is the worst item still in the set, which is the one to evict. A new item enters only if it beats the root. Python's heapq module only provides a min-heap, so the usual trick is to reverse the comparison with a small wrapper. Negating keys only works for numbers.

import heapq
from itertools import islice

class _Rev:
    """Reverses ordering so heapq's min-heap behaves as a max-heap."""
    __slots__ = ("item",)
    def __init__(self, item): self.item = item
    def __lt__(self, other): return other.item < self.item

def partial_sort_heap(iterable, k):
    """Return the k smallest items in ascending order. O(n log k) time, O(k) memory."""
    if k <= 0:
        return []
    it = iter(iterable)
    heap = [_Rev(x) for x in islice(it, k)]
    heapq.heapify(heap)                      # O(k)
    for x in it:
        if x < heap[0].item:                 # beats the current k-th smallest
            heapq.heapreplace(heap, _Rev(x)) # pop root and push x in one sift
    return sorted(r.item for r in heap)      # O(k log k)

This version accepts any iterator, so it works on a file or a network stream that never fits in memory. It uses heapreplace instead of a separate pop and push, which saves one sift. The strict < in the test means that when items tie with the current k-th item, the first one seen is kept. If you need the result to be stable, use (key, sequence_number) pairs as the items. The standard library already provides this routine: heapq.nsmallest(k, iterable, key=...) is documented as equivalent to sorted(iterable, key=key)[:k].

Selection, then sort the prefix

When the whole array is in memory and you are allowed to reorder it, selection is the asymptotically better route. Partition around a random pivot, keep only the side that contains position k-1, and repeat. Once the pivot lands at index k-1, the prefix holds the k smallest items and only needs a sort.

import random

def partition(a, lo, hi):
    r = random.randint(lo, hi)           # random pivot defeats sorted-input traps
    a[r], a[hi] = a[hi], a[r]
    pivot, i = a[hi], lo
    for j in range(lo, hi):
        if a[j] < pivot:
            a[i], a[j] = a[j], a[i]
            i += 1
    a[i], a[hi] = a[hi], a[i]
    return i

def select_then_sort(a, k):
    """In place: a[0:k] becomes the k smallest, ascending. Expected O(n + k log k)."""
    lo, hi = 0, len(a) - 1
    while lo < hi:
        p = partition(a, lo, hi)
        if p == k - 1:
            break
        if p < k - 1:
            lo = p + 1
        else:
            hi = p - 1
    a[:k] = sorted(a[:k])

def partial_quicksort(a, lo, hi, k):
    """Quicksort that skips every subarray lying entirely at or beyond index k."""
    while lo < hi:
        p = partition(a, lo, hi)
        if p < k - 1:
            partial_quicksort(a, lo, p - 1, k)   # left part is inside the prefix: sort it
            lo = p + 1                           # right part straddles k: keep going
        else:
            hi = p - 1                           # right part is all beyond k: skip it

Partial quicksort, analysed by Conrado Martinez in 2004, combines the two steps. It is ordinary quicksort, except that it never recurses into a part that lies entirely past index k-1. Its expected cost is O(n + k log k), the same order as select-then-sort, and it needs one routine instead of two. libstdc++'s nth_element and NumPy's default kind='introselect' use introselect. It starts like quickselect and switches to a fallback when recursion gets too deep, which guards against adversarial inputs. Check your own library's worst case. If you write your own selection, add a depth limit, or at least use random pivots.

What the libraries give you

You rarely need to write any of this yourself. These are the standard tools and what they actually guarantee:

// C++: k slowest requests, slowest first. Heap-based; about N*log(k) comparisons; NOT stable.
std::partial_sort(reqs.begin(), reqs.begin() + k, reqs.end(),
                  [](const Req& a, const Req& b) { return a.latency_us > b.latency_us; });

// C++: same result via selection, then sort only the prefix. nth_element is linear on average.
std::nth_element(reqs.begin(), reqs.begin() + (k - 1), reqs.end(), by_latency_desc);
std::sort(reqs.begin(), reqs.begin() + k, by_latency_desc);

// Rust (k >= 1 and k <= v.len()): select, then sort the prefix.
v.select_nth_unstable(k - 1);
v[..k].sort_unstable();

# NumPy: indices of the k smallest scores, in order, without sorting all n.
idx = np.argpartition(scores, k - 1)[:k]   # introselect by default, linear on average
idx = idx[np.argsort(scores[idx])]         # order only the k winners

SQL does the same thing for you. In PostgreSQL, ORDER BY latency DESC LIMIT 20 on an unindexed column shows Sort Method: top-N heapsort in EXPLAIN ANALYZE when the bounded heap fits in work_mem. If you see a full external sort instead, check whether the LIMIT reaches the sort node or a subquery hides it. An index on the sort key avoids the problem entirely, because the first k rows of an ordered index scan are the answer.

Worked example: counting comparisons

To compare the methods without depending on a particular machine, count comparisons. Each input item was wrapped in a class whose __lt__ increments a counter. The heap and the two selection methods ran with n = 100,000 and k = 100 on 20 random seeds, and the heap and quickselect also ran on reverse-sorted input. The full sort was measured separately. Results were checked against sorted(vals)[:k]. These are the measured comparisons per input element:

Three lessons come out of these numbers. First, on random input the heap does barely more than one comparison per element, because almost every item loses to the root. Second, quickselect has high variance: one unlucky early pivot costs another full pass. Third, adversarial order changes the result. On descending input every new item beats the root, so the heap pays a full sift each time, about log2(100), or 7 levels. The randomised selection does not care about input order. Counting comparisons does not capture memory traffic, though. The heap reads the input once, sequentially, and its k items stay in L1 cache. Selection makes several passes and swaps elements across the whole array. In practice the heap is often faster in wall-clock time even when the comparison counts are equal, so measure on your own data before deciding.

Sharded and streaming partial sort

Partial sort composes well across machines. The global top k is always contained in the union of each shard's local top k. So each shard returns its k best, already sorted, and the coordinator runs a k-way merge that stops after k outputs. That costs O(k log s) for s shards. The k-way heap merge is the same algorithm. This is how sharded search and log-analytics systems answer top-N queries, and it explains why deep pagination is expensive. Page p of size m forces every shard to return p*m candidates. The coordinator then throws almost all of them away. Cursor-based pagination avoids this by remembering the last key and asking each shard for items strictly after it.

For a sliding window, such as the 20 slowest requests in the last five minutes, a heap alone is not enough because items have to expire. Use an indexed heap with removal, or recompute per time bucket and merge the buckets.

Failure modes

• Wrong heap direction. A min-heap of size k keeps the k largest items. This is the most common bug, and it passes tests that only check the result length.
• Off-by-one in selection. nth_element(begin, begin + k, end) puts the (k+1)-th smallest item at index k, while the prefix [0, k) is still the k smallest. Mixing that convention with k - 1 elsewhere drops or adds an item. Pick one convention and test k = 1 and k = n.
• Assuming stability. std::partial_sort and selection-based methods are not stable. If equal keys must keep their input order, for example in pagination, add a unique tie-breaker to the comparator.
• Inconsistent comparators. NaN keys break strict weak ordering. C++ may then produce garbage or read out of bounds, and NumPy sorts NaNs to the end, which may not be what you want. Filter or map NaNs before the call.
• k larger than n or k = 0. Clamp k before the call. Several APIs treat an out-of-range k as an error or undefined behaviour.
• Quadratic quickselect. A first-element or last-element pivot on already sorted data makes hand-written quickselect O(n^2). Use random pivots or introselect.

Trade-offs

Heap versus selection. Choose the heap for streams, for inputs you may not modify, and when k is small. Choose selection when the data is already in a mutable array and k is a large fraction of n, because n log k grows with k but n + k log k grows more slowly until k nears n. Above roughly k = n/2, a full sort is simpler and close in cost. Copy versus in place. Selection reorders the input, so if callers need the original order, you pay O(n) for a copy, and the heap usually wins. Exact versus approximate. If the requirement is only a set of good candidates, as in approximate nearest-neighbour search, you can skip exact partial sorting and use a cheaper threshold or a sampled cut-off.

What to do next

1. Search your code for sort( followed by [:k], head( or LIMIT, and replace each with heapq.nsmallest, partial_sort or argpartition where n is large.
2. Decide whether stability matters for each call site and add a tie-breaker key if it does.
3. Write property tests that compare against the full sort for k in {0, 1, n-1, n}, with duplicates and with sorted and reverse-sorted input.
4. Run EXPLAIN ANALYZE on your top-N queries and confirm the plan shows a top-N heapsort or an index scan.
5. For sharded data, return the sorted local top k from each shard, merge at the coordinator, and move deep pagination to cursors.
6. Keep learning with heap operations and k-way merging.