1. Top-K Retrieval Candidates

Problem

Fin's retrieval stage returns roughly 40 scored candidate passages per customer query. Before generation, it needs only the top-k highest-scoring passages, and this runs on every single query in production, so a full O(n log n) sort of the whole batch is wasteful.

Given the unsorted (doc_id, score) pairs, return only the k best. The result must be sorted descending by score, with ties broken by doc_id ascending, and the whole thing must run in O(n log k) time using a heap — not a full sort.

Source: src/1_top_k_retrieval_candidates.py

def top_k_candidates(candidates: list[tuple[str, float]], k: int) -> list[tuple[str, float]]:
    ...

Examples:
>>> top_k_candidates([("a", 0.9), ("b", 0.3), ("c", 0.7), ("d", 0.5)], 2)
[("a", 0.9), ("c", 0.7)]

>>> top_k_candidates([("a", 0.5), ("b", 0.5), ("c", 0.9)], 2)
[("c", 0.9), ("a", 0.5)]

Step-by-Step Approach

1. Recognize that you don't need a full sort — you need only the k largest elements out of n, so O(n log k) beats O(n log n) whenever k is much smaller than n.
2. Maintain a min-heap bounded to size k. The heap's root is always the "worst of the best": the smallest score among the k candidates currently kept.
3. For each candidate: if the heap has fewer than k items, push it unconditionally. Otherwise compare it to the heap's smallest kept item; if the candidate is better, evict the smallest and push the candidate in its place.
4. Because ties must break deterministically (score tie → doc_id ascending), encode a composite comparison key so the heap orders correctly even on ties, instead of comparing raw scores alone.
5. After processing every candidate, the heap holds exactly the top-k, but not in the required output order — pull them out and do one small final sort (O(k log k)) to produce the descending, tie-broken order the problem asks for.
6. Handle edge cases: k = 0 leaves the heap empty, empty input candidates give an empty result, and k larger than n just means the heap ends up holding all n elements.

The key insight is that the heap size stays bounded at k throughout, so each element costs O(log k) to process instead of O(log n) — you never pay to fully order the elements you're about to discard.

Reference solution

import heapq


def top_k_candidates(candidates: list[tuple[str, float]], k: int) -> list[tuple[str, float]]:
    # bounded min-heap of size k keyed on (score, reversed doc_id), O(n log k) time, O(k) space
    heap: list[tuple[tuple[float, tuple[int, ...]], str, float]] = []
    for doc_id, score in candidates:
        # negate ord() to flip tie-break to ascending
        tie_key = tuple(-ord(c) for c in doc_id)
        # tuple compares lexicographically: score, then tie_key
        entry_key = (score, tie_key)
        if len(heap) < k:
            # grow heap to size k
            heapq.heappush(heap, (entry_key, doc_id, score))
        elif k > 0 and entry_key > heap[0][0]:
            # evict worst-of-best in one op
            heapq.heapreplace(heap, (entry_key, doc_id, score))
    # final sort: score desc, doc_id asc
    ordered = sorted(heap, key=lambda item: (-item[2], item[1]))
    return [(doc_id, score) for _, doc_id, score in ordered]

Key Functions & Tricks

• heapq.heappush(heap, item) — push, sift up, O(log n); grows heap to size k
• heapq.heapreplace(heap, item) — atomic pop+push, O(log n), evicts worst-of-best
• Bounded min-heap for top-k max query — min-heap capped at size k tracks k largest
• tuple(-ord(c) for c in doc_id) — negate ord() per char to flip tie-break direction
• entry_key = (score, tie_key) — tuple packs primary + tie-break for lexicographic compare
• sorted(heap, key=lambda item: (-item[2], item[1])) — final cleanup: score desc, doc_id asc

How to Recognize This Pattern