beam: k-way merge of sorted streams with stability + tests

wizard-work/beam/k_way_merge.py

@@ -0,0 +1,49 @@
+"""k_way_merge.py — merge K sorted (ascending) iterables into one sorted stream.
+
+Deterministic, stable (ties broken by input order), handles empty inputs,
+single-element inputs, and arbitrary iterables. O(n log K) where n is total
+elements and K is number of inputs.
+"""
+from __future__ import annotations
+import heapq
+from typing import Iterable, Iterator, TypeVar
+
+T = TypeVar("T")
+
+
+def merge_k_sorted(streams: Iterable[Iterable[T]]) -> Iterator[T]:
+    """Merge K sorted ascending streams into one sorted ascending iterator.
+
+    Stability: when two elements compare equal, the one from the lower-index
+    stream is emitted first. Empty streams are skipped. Inputs need not be
+    lists; any iterable is consumed lazily per stream.
+    """
+    # Materialise into a list so we can index for tie-breaking stability.
+    streams = list(streams)
+    # heap entries: (value, stream_index, seq, iterator)
+    # seq is a global counter to make entries comparable when values tie and
+    # stream_index alone is not enough (same stream can't have two live entries
+    # at once, but including seq keeps heap entries totally ordered).
+    heap: list = []
+    iters = [iter(s) for s in streams]
+    seq = 0
+    for i, it in enumerate(iters):
+        try:
+            v = next(it)
+        except StopIteration:
+            continue
+        heapq.heappush(heap, (v, i, seq, it))
+        seq += 1
+    while heap:
+        v, i, _seq, it = heapq.heappop(heap)
+        yield v
+        try:
+            nv = next(it)
+        except StopIteration:
+            continue
+        heapq.heappush(heap, (nv, i, seq, it))
+        seq += 1
+
+
+def merge_k_sorted_list(streams: Iterable[Iterable[T]]) -> list[T]:
+    return list(merge_k_sorted(streams))

wizard-work/beam/test_k_way_merge.py

@@ -0,0 +1,75 @@
+import random
+from k_way_merge import merge_k_sorted, merge_k_sorted_list
+
+
+def test_empty():
+    assert merge_k_sorted_list([]) == []
+    assert merge_k_sorted_list([[], [], []]) == []
+
+
+def test_single():
+    assert merge_k_sorted_list([[1, 2, 3]]) == [1, 2, 3]
+
+
+def test_basic():
+    a = [1, 4, 7]
+    b = [2, 5, 8]
+    c = [3, 6, 9]
+    assert merge_k_sorted_list([a, b, c]) == list(range(1, 10))
+
+
+def test_duplicates_stability():
+    # equal values across streams: lower index first
+    a = [5, 5]
+    b = [5, 5]
+    out = merge_k_sorted_list([a, b])
+    assert out == [5, 5, 5, 5]
+
+
+def test_empty_streams_mixed():
+    assert merge_k_sorted_list([[], [1], [], [2, 3]]) == [1, 2, 3]
+
+
+def test_iterables_not_lists():
+    def gen(n):
+        for i in range(n):
+            yield i * 2
+    assert merge_k_sorted_list([gen(3), [1, 3, 5]]) == [0, 1, 2, 3, 4, 5]
+
+
+def test_large_random_vs_sorted():
+    rng = random.Random(42)
+    streams = []
+    for _ in range(7):
+        s = sorted(rng.randint(0, 1000) for _ in range(200))
+        streams.append(s)
+    out = merge_k_sorted_list(streams)
+    expected = sorted(x for s in streams for x in s)
+    assert out == expected
+
+
+def test_total_count():
+    streams = [[1, 2], [3], [], [4, 5, 6]]
+    assert len(merge_k_sorted_list(streams)) == 6
+
+
+def test_single_elements():
+    assert merge_k_sorted_list([[9], [1], [5]]) == [1, 5, 9]
+
+
+def test_descending_not_assumed():
+    # inputs assumed sorted ascending; verify we don't sort inputs
+    assert merge_k_sorted_list([[1, 2], [3, 4]]) == [1, 2, 3, 4]
+
+
+def test_iterator_lazy():
+    it = merge_k_sorted([[1, 2], [3]])
+    assert next(it) == 1
+    assert next(it) == 2
+    assert next(it) == 3
+    try:
+        next(it)
+    except StopIteration:
+        pass
+    else:
+        raise AssertionError("expected StopIteration")