2. Merge K Ranked Result Lists

Problem

Fin's retrieval draws from three distinct source categories and synthesizes an answer by combining relevant snippets across multiple sources. Each source category returns its own already-ranked list of (doc_id, score) snippets.

Before generation, these k independently-ranked lists need to be merged into one globally ranked list. Each inner list is already sorted descending by score; merge all k lists into one list sorted descending by score in O(N log k) time (N = total elements across all lists) using a heap-based k-way merge — not concatenate-then-sort.

Source: src/2_merge_ranked_results_k_sources.py

def merge_k_ranked(lists: list[list[tuple[str, float]]]) -> list[tuple[str, float]]:
    ...

Examples:
>>> merge_k_ranked([[("a", 0.9), ("b", 0.5)], [("c", 0.8), ("d", 0.2)]])
[("a", 0.9), ("c", 0.8), ("b", 0.5), ("d", 0.2)]

>>> merge_k_ranked([])
[]

>>> merge_k_ranked([[], [("a", 0.5)], []])
[("a", 0.5)]

Step-by-Step Approach

1. Recognize this as "merge sorted streams" generalized from 2 lists to k lists — the classic k-way merge.
2. Seed a min-heap with the head element of each of the k lists (negating scores so a min-heap behaves like a max-heap for "highest score first"), tagging each heap entry with which list and which index it came from so that specific list can be advanced later.
3. Repeatedly pop the globally-best head element from the heap and append it to the output.
4. Immediately after popping an element from list i at index idx, push list i's next element (idx + 1) onto the heap if it exists — this "advance the just-consumed lane" step is what keeps the heap bounded to at most k live entries at all times.
5. Remember to negate the popped score back to its original sign before appending to the result, since it was negated only to simulate max-heap ordering with Python's min-heap-only heapq.
6. Handle edge cases: skip seeding empty inner lists, and an empty outer list means the heap starts empty, the loop never runs, and the result is simply [].

The key insight is that at any moment the heap holds exactly one live candidate per non-exhausted list, so it never grows beyond k items — giving O(N log k) total work instead of the O(N log N) a concatenate-then-sort approach would cost.

Reference solution

import heapq


def merge_k_ranked(lists: list[list[tuple[str, float]]]) -> list[tuple[str, float]]:
    # heap-based k-way merge, one head element per list, O(N log k) time, O(k) space
    heap: list[tuple[float, str, int, int]] = []
    # track index i to know which lane to advance later
    for i, lst in enumerate(lists):
        if lst:
            doc_id, score = lst[0]
            # negate score so min-heap yields max first
            heapq.heappush(heap, (-score, doc_id, i, 0))
    result: list[tuple[str, float]] = []
    while heap:
        neg_score, doc_id, i, idx = heapq.heappop(heap)
        # flip back to true score
        result.append((doc_id, -neg_score))
        if idx + 1 < len(lists[i]):
            next_doc_id, next_score = lists[i][idx + 1]
            # advance this lane, keep heap <= k
            heapq.heappush(heap, (-next_score, next_doc_id, i, idx + 1))
    return result

Key Functions & Tricks

• heapq.heappush(heap, item) — push + sift, O(log k) since heap holds at most k entries
• heapq.heappop(heap) — pop smallest negated-score entry, O(log k), yields global max
• Negate-the-score trick for a max-heap — store -score so heap root is highest actual score
• enumerate(lists) — pairs index i with each list to seed heap and track advance later
• 4-tuple heap entries (-score, doc_id, i, idx) — bookkeeping for tie-break and lane advance
• "Advance the just-consumed lane" step — push list i's next element after popping its head

How to Recognize This Pattern