Merge ordered streams

Question: Merge ordered streams

Companies: Sierra, Amazon

Difficulty: Medium

Problem statement

Merge several ordered input streams into one ordered output without loading any stream fully into memory. Reported on 1point3acres' Sierra bank; reconstructed in our own words.

Clarify the setup

You get k input streams, each already sorted, and must produce one sorted output stream. A stream can be huge or endless (an event log, a paginated API), so nothing may be read fully into memory. Ask:

• Interface. Python iterators in, a generator out. Each next() on the output should pull only what it needs from the inputs.
• Ties. Equal values from different streams: output order by stream index (stable), so results are deterministic.
• Duplicates. Keep them (a merge) or drop them (a union)? Make it an option.
• Sort key. Records such as (timestamp, event) sorted by a key function.
• Bad input. A stream that goes backwards breaks the output order silently. Detect it and raise.

Approach

Brute force: read everything, sort: O(N log N) time and O(N) memory. Ruled out by the streaming requirement.

Pairwise merging (merge stream 1 with 2, the result with 3, …) streams fine but costs O(N·k), since early elements pass through up to k − 1 merges.

Min-heap of stream heads: O(N log k) time, O(k) memory.

1. Pull the first element of each non-empty stream and push (key, stream_index, value) onto a min-heap. The stream index breaks ties, which gives stability and also means Python never compares two values directly (values may not be comparable).
2. Pop the smallest, yield its value, and push the next element of the same stream, if any.
3. Stop when the heap is empty.

Laziness matters: the generator pulls a stream's next element only after yielding that stream's previous one, so an endless stream merged with finite ones works, and a consumer that stops early never drains the inputs.

Dropping duplicates is one comparison against the last yielded key. Detecting a backwards stream is one comparison against that stream's previous key.

Code

import heapq
from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple


class StreamOrderError(ValueError):
    pass


def merge_streams(
    streams: Iterable[Iterable[Any]],
    key: Optional[Callable[[Any], Any]] = None,
    unique: bool = False,
) -> Iterator[Any]:
    keyf = key or (lambda v: v)
    iters = [iter(s) for s in streams]
    heap: List[Tuple[Any, int, Any]] = []

    def push(i: int) -> None:
        for value in iters[i]:            # at most one element
            heapq.heappush(heap, (keyf(value), i, value))
            return

    for i in range(len(iters)):
        push(i)

    last_by_stream: List[Any] = [None] * len(iters)
    seen_by_stream = [False] * len(iters)
    have_last, last_key = False, None
    while heap:
        k, i, value = heapq.heappop(heap)
        if seen_by_stream[i] and k < last_by_stream[i]:
            raise StreamOrderError(f"stream {i} went backwards: {k!r} after {last_by_stream[i]!r}")
        last_by_stream[i], seen_by_stream[i] = k, True
        if not (unique and have_last and k == last_key):
            yield value
        have_last, last_key = True, k
        push(i)


def merge_two(a: Iterable[Any], b: Iterable[Any]) -> Iterator[Any]:
    """The two-stream case without a heap, for comparison."""
    ia, ib = iter(a), iter(b)
    sentinel = object()
    x, y = next(ia, sentinel), next(ib, sentinel)
    while x is not sentinel and y is not sentinel:
        if y < x:
            yield y
            y = next(ib, sentinel)
        else:                              # ties take from a first: stable
            yield x
            x = next(ia, sentinel)
    while x is not sentinel:
        yield x
        x = next(ia, sentinel)
    while y is not sentinel:
        yield y
        y = next(ib, sentinel)

The backwards check compares the new key with the same stream's previous key at pop time. A stream that goes backwards is caught when its out-of-order element reaches the top of the heap, which is before anything later than it could be yielded out of order from that stream.

Complexity

• Time: O(N log k) for N total elements; each element is pushed and popped once.
• Memory: O(k), one head per stream.
• With k = 1,000 streams and N = 10^9 events, that is about 10^10 heap operations' worth of comparisons (log2 of 1,000 is 10): minutes in C, much longer in pure Python, which is why large merges are sharded (follow-up 3).

Test cases

import itertools as _it
import random as _random


def test_basic_merge_and_empty_streams():
    assert list(merge_streams([[1, 4, 7], [], [2, 5], [3, 6, 8, 9]])) == list(range(1, 10))
    assert list(merge_streams([])) == []
    assert list(merge_streams([[], []])) == []


def test_matches_sorted_on_random_input():
    rng = _random.Random(2)
    for _ in range(100):
        streams = [sorted(rng.randint(0, 50) for _ in range(rng.randint(0, 20))) for _ in range(rng.randint(1, 8))]
        assert list(merge_streams(streams)) == sorted(_it.chain(*streams))


def test_stable_ties_and_key():
    a = [(1, "a1"), (2, "a2")]
    b = [(1, "b1"), (2, "b2")]
    out = list(merge_streams([a, b], key=lambda r: r[0]))
    assert out == [(1, "a1"), (1, "b1"), (2, "a2"), (2, "b2")]
    assert list(merge_two([1, 3], [1, 2])) == [1, 1, 2, 3]

def test_lazy_with_infinite_stream():
    pulled = []

    def evens():
        n = 0
        while True:
            pulled.append(n)
            yield n
            n += 2

    out = list(_it.islice(merge_streams([evens(), [1, 3, 5]]), 6))
    assert out == [0, 1, 2, 3, 4, 5]
    assert len(pulled) <= 4                 # only pulled what it needed


def test_unique_and_backwards_detection():
    assert list(merge_streams([[1, 2, 2, 3], [2, 3, 4]], unique=True)) == [1, 2, 3, 4]
    try:
        list(merge_streams([[1, 5, 2], [3, 4]]))
        assert False
    except StreamOrderError:
        pass


def test_uncomparable_values_with_equal_keys():
    class Event:
        def __init__(self, t):
            self.t = t

    out = list(merge_streams([[Event(1), Event(2)], [Event(1)]], key=lambda e: e.t))
    assert [e.t for e in out] == [1, 1, 2]

Follow-ups

1. Merging pages from APIs. Each stream is a paginated endpoint. Wrap each in a generator that fetches the next page when its buffer runs out; the merge stays unchanged and only fetches pages it needs.
2. Async streams. With async iterators, the heap logic is the same, but pulling the next element is awaited. Prefetch one page per stream concurrently so a slow stream does not stall the others more than necessary.
3. Huge merges. For thousands of streams or files (external sort), merge in a tree: groups of 100 streams merge into intermediate streams, which merge again. It caps open file handles and parallelizes across machines.
4. Endless streams and watermarks. With live event streams, the merge cannot emit time t until every stream has shown it is past t. A stream that goes quiet blocks the output; streams send periodic watermarks ("nothing older than t will come"), and the merge emits up to the minimum watermark.
5. Top-k or windowed output. Stop early once enough results are produced (the lazy design already supports this), or merge only within a time window.

Interviewer expectations

Mid-level (L4)

• Uses a min-heap of stream heads and gets the right output on normal inputs, with empty streams handled.
• States O(N log k) time and O(k) memory.
• Tests against sorting the concatenated input.

Senior (L5)

• Writes it as a lazy generator that never reads a stream fully, and shows it works with an infinite stream.
• Uses the stream index as a tie-breaker for stability and to avoid comparing values; supports a key function and de-duplication.
• Detects out-of-order input.

Staff (L6+)

• Extends to async and paginated sources with prefetching, tree merges for very large k, and watermarks for live streams.
• Discusses back-pressure and what happens when one source is slow or stalls.
• Picks the design from the data's sizes and the consumer's needs.