50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""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))
|