Top-k selection

Deep dive

How it works

In: a stream or array of n scored items, k · Out: the k best items, sorted · State: a min-heap of k items (or a buffer of 2k and a threshold) · Cost: heap O(n log k) worst case, about O(n) on random input; buffered quickselect O(n) expected

Bounded heap (streaming):

1. Push the first k items into a min-heap.
2. For each later item x: if x > root, replace the root and sift down (log2 k steps); else drop x.
3. At the end, sort the heap (k log k).

Buffered selection (array or stream, no per-item heap work):

1. Append items that beat the current threshold to a buffer of capacity 2k.
2. When full, run quickselect on the buffer to keep the best k, and set threshold = the minimum of those k.
3. At the end, select once more and sort the k survivors. Each prune costs O(k) and happens at most once per k accepted items, so the total is O(n) expected.

Quickselect: partition around a pivot as in quicksort, then recurse only into the side containing position k. Introselect falls back to median of medians to guarantee O(n).

stream: 7 3 9 1 8 2 6        k = 3
heap  : [3 7 9] min 3 -> 1 drop -> 8 replaces 3: [7 8 9] -> 2 drop -> 6 drop
result: 9 8 7

Code

use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// k largest items from a stream, O(n log k) worst case, O(k) memory. Result sorted descending.
pub fn top_k_heap<T: Ord, I: IntoIterator<Item = T>>(items: I, k: usize) -> Vec<T> {
    if k == 0 { return Vec::new(); }
    let mut heap = BinaryHeap::with_capacity(k); // min-heap of the current best k
    for x in items {
        if heap.len() < k {
            heap.push(Reverse(x));
        } else if let Some(mut min) = heap.peek_mut() {
            if x > min.0 { *min = Reverse(x); } // replace-top: one sift-down; usually skipped
        }
    }
    let mut v: Vec<T> = heap.into_iter().map(|Reverse(x)| x).collect();
    v.sort_unstable_by(|a, b| b.cmp(a));
    v
}

/// k largest via a 2k buffer pruned with quickselect: O(n) expected, branch-light filter.
pub fn top_k_select<T: Ord + Copy, I: IntoIterator<Item = T>>(items: I, k: usize) -> Vec<T> {
    if k == 0 { return Vec::new(); }
    let mut buf: Vec<T> = Vec::with_capacity(2 * k);
    let mut threshold: Option<T> = None; // k-th largest seen so far: smaller items can never qualify
    for x in items {
        if threshold.map_or(true, |t| x > t) {
            buf.push(x);
            if buf.len() == 2 * k {
                buf.select_nth_unstable_by(k - 1, |a, b| b.cmp(a)); // best k in buf[..k]
                buf.truncate(k);
                threshold = buf.iter().min().copied();
            }
        }
    }
    if buf.len() > k { buf.select_nth_unstable_by(k - 1, |a, b| b.cmp(a)); buf.truncate(k); }
    buf.sort_unstable_by(|a, b| b.cmp(a));
    buf
}

/// Merge per-shard top-k lists (each sorted descending) into the global top-k.
pub fn merge_top_k<T: Ord + Copy>(shards: &[Vec<T>], k: usize) -> Vec<T> {
    let mut heap: BinaryHeap<(T, usize, usize)> = shards.iter().enumerate()
        .filter(|(_, s)| !s.is_empty()).map(|(i, s)| (s[0], i, 0)).collect();
    let mut out = Vec::with_capacity(k);
    while out.len() < k {
        let Some((x, i, j)) = heap.pop() else { break };
        out.push(x);
        if let Some(&y) = shards[i].get(j + 1) { heap.push((y, i, j + 1)); }
    }
    out
}

Complexity and constants

Measured (code above, 10^7 random u64, one Apple M-series core, ms):

• Heap replacements on random order: about k·ln(n/k); each costs order of log2(k) cache-resident compares. On ascending input every item replaces the root: n·log2 k.
• select_nth_unstable on the whole array is O(n) (about 1 ns/item here) but needs the array in memory and mutable.
• Memory: k items for the heap, 2k for the buffered version.

Efficiency

• Cost: measured with the code above (10^7 random u64, one Apple M-series core): heap top-10 in 5.1 ms (about 0.5 ns/item), buffered select 3.8 ms, full sort 114 ms. At k = 10^6 the heap needs 195 ms and buffered select 35 ms; on ascending input with k = 100 the heap needs 123 ms (every item replaces the root) and buffered select 31 ms. Memory: k items (2k buffered). Bound: one predictable compare per item (memory bandwidth), then sift-down work once many items qualify.
• Levers:
  • threshold compare before any heap call — a single predictable branch per item; most items drop — random or score-skewed input
  • buffered quickselect instead of a heap — 3-4 ns/item where the heap hits 12-14 (measured above) — large k or ascending input
  • SIMD filter against the threshold, then compress survivors — fewer scalar branches (order of) — dense score arrays
  • early termination by upper bound (WAND, MaxScore, block-max; Ding & Suel 2011) — skips most posting lists — index-backed ranking
  • 4-ary heap (d-ary heap) — shallower sift-down, fewer cache misses (order of) — large k
• Naive → standard → state of the art:
  • full sort, take k · 114 ms for 10^7 u64 (measured above)
  • bounded min-heap · 5.1 ms at k = 10, 195 ms at k = 10^6 (measured above)
  • threshold filter plus buffered quickselect, bound-based skipping in indexes · 3.8 ms at k = 10, 35 ms at k = 10^6 (measured above)

Implementation notes

• Replace-top, not pop + push: BinaryHeap::peek_mut and assigning through it does one sift-down when the guard drops.
• Compare with the threshold before anything else: the hot path is a single predictable branch per item; keep the threshold in a register (as threshold above), not behind the heap API.
• Vectorize the filter: compare 8-16 scores against the threshold with SIMD, compress the survivors, then heap or buffer them. Search engines' top-k collectors work this way.
• Ties: make keys total and deterministic, e.g. (score, doc_id), so shards and retries return the same set.
• Floats: use total_cmp or a key like key_f64 from radix sort; NaN breaks PartialOrd-based heaps.
• Shards: each returns its top k sorted; merge with a k-way heap as merge_top_k does, popping only k times.
• Early termination in indexes: if items arrive in descending order of an upper bound (WAND, MaxScore, block-max indexes), stop once the bound falls below the threshold. This is how BM25 top-k skips most of the posting lists.
• d-ary heaps (4-ary) reduce sift-down depth and cache misses for large k.

Pitfalls

• Using a max-heap of all n items and popping k: O(n + k log n) time but O(n) memory, which defeats streaming.
• Heap direction confusion: top-k largest needs a min-heap (Reverse in Rust's max-heap BinaryHeap).
• Pathological input orders (ascending scores, time-sorted data) turn the heap into n log k. Use the buffered version or shuffle chunks.
• Merging shard results by asking every shard for k/shards items: wrong. Each shard must return k (or use a global threshold protocol).
• Large k with large records: move indices or (key, id) pairs, not whole rows.

Variants and alternatives

In production

Standard libraries: C++ std::partial_sort (heap-based) and std::nth_element (introselect), Rust select_nth_unstable. Search engines: Lucene TopScoreDocCollector with block-max WAND. Databases: PostgreSQL's bounded heap sort for ORDER BY … LIMIT (EXPLAIN shows "top-N heapsort"), ClickHouse, DuckDB and DataFusion top-N operators. Vector search: the top-k operators of FAISS and other engines.

Implementations

Sources

• C. A. R. Hoare, Algorithm 65: Find, Communications of the ACM, 1961. paper
• M. Blum, R. W. Floyd, V. Pratt, R. L. Rivest, R. E. Tarjan, Time Bounds for Selection, Journal of Computer and System Sciences, 1973. paper
• R. W. Floyd, R. L. Rivest, Expected Time Bounds for Selection, Communications of the ACM, 1975. paper
• D. R. Musser, Introspective Sorting and Selection Algorithms, Software: Practice and Experience, 1997.
• S. Ding, T. Suel, Faster Top-k Document Retrieval Using Block-Max Indexes, SIGIR, 2011. paper
• Rust standard library, slice::select_nth_unstable. docs