~/problems / Iterators & parsers / Iterators and generators

Merge K sorted streams

medium 2 levels ~40 min Citadel

Level 1 A merged stream

A stream is any object with three methods:

  • has_next() -> bool
  • peek(): the next value, without consuming it
  • next(): return the next value and advance

Calling peek() or next() on an exhausted stream raises StopIteration.

Write two stream classes:

  1. ListStream(values), a stream over a Python list.
  2. MergedStream(streams), which takes a list of streams, each already sorted in non-decreasing order, and is itself a stream that yields all their values in non-decreasing order. Equal values from different input streams come out in input order (the lower index first). Also make it iterable: __iter__ returns self and __next__ behaves like next(), so list(MergedStream(...)) works.

Some rules the tests check:

  • The inputs can be any stream objects, not just your ListStream, and may be infinite. Never read an input to the end up front; only look at each input's current head.
  • next() on an input consumes a value. Call it only when that value is the one you're about to emit. Don't call next() just to look.
  • With k inputs and N values in total, the full merge must take O(N log k). Scanning all k heads for every value is too slow for the tests (k = 1000).
m = MergedStream([ListStream([1, 4, 7]), ListStream([2, 5]), ListStream([]), ListStream([3, 4])])
m.peek()      # 1
list(m)       # [1, 2, 3, 4, 4, 5, 7]
m.has_next()  # False

Level 2 unlocks when level 1 passes.

Topic: Iterators and generators. Resumable/serializable iterators, merging streams, lazy pipelines.

0:00
Ctrl ' run · Ctrl ↵ submit
esc