Problem 90 · Data Engineering & Streaming

Merge K Sorted Iterators

MediumClassic DSA & systems

01Problem

Write a generator that lazily merges k sorted (possibly infinite) iterators into one sorted stream, with an optional key and stable tie-breaking.

02Key insight

A min-heap of (key, source_index, value, iterator) gives O(log k) per element. The source index both makes ties stable and prevents Python from ever comparing iterators or un-orderable values; heapreplace avoids a pop+push.

03Solution

solution.pypython24 lines

import heapq
from typing import Any, Callable, Iterable, Iterator, Optional, TypeVar

T = TypeVar("T")
_DONE = object()


def merge_sorted(*iterables: Iterable[T], key: Optional[Callable[[T], Any]] = None) -> Iterator[T]:
    keyfn = key or (lambda x: x)
    heap: list[tuple[Any, int, T, Iterator[T]]] = []
    for idx, iterable in enumerate(iterables):
        it = iter(iterable)
        first = next(it, _DONE)
        if first is not _DONE:
            heap.append((keyfn(first), idx, first, it))
    heapq.heapify(heap)
    while heap:
        _, idx, value, it = heap[0]
        yield value
        nxt = next(it, _DONE)
        if nxt is _DONE:
            heapq.heappop(heap)
        else:
            heapq.heapreplace(heap, (keyfn(nxt), idx, nxt, it))

Try it first.

TimeO(n log k)

SpaceO(k)

tests.py · 3 passing

tests.pypython6 lines

import itertools
assert list(merge_sorted([1, 4, 7], [2, 5], [], [0, 3, 6, 9])) == [0, 1, 2, 3, 4, 5, 6, 7, 9]
evens, odds = itertools.count(0, 2), itertools.count(1, 2)
assert list(itertools.islice(merge_sorted(evens, odds), 6)) == [0, 1, 2, 3, 4, 5]
rows = merge_sorted([(1, "a"), (2, "a")], [(1, "b"), (2, "b")], key=lambda r: r[0])
assert list(rows) == [(1, "a"), (1, "b"), (2, "a"), (2, "b")]  # stable across sources

04Staff-level follow-up

Sources are remote shards with network latency - what changes? Prefetch a buffered block per source (async or threads) so the heap never stalls on one slow shard, and bound total buffer memory; this is the coordinator side of a distributed ORDER BY ... LIMIT.