merge_streams_better.py

from collections import namedtuple
IndexValue = namedtuple("IndexValue", "index value")

def merge_streams_better(streams):
    if len(streams) == 1:
        return list(streams[0])

    merged_list = []
    stream_pairs = [
        IndexValue(index, next(value)) for index, value in enumerate(streams)
    ]
    while any(pair.value is not None for pair in stream_pairs):
        smallest = min(
            [pair for pair in stream_pairs if pair.value is not None],
            key=lambda x: x.value,
        )
        merged_list.append(smallest.value)
        try:
            stream_pairs[smallest.index] = IndexValue(
                smallest.index, next(streams[smallest.index])
            )
        except StopIteration:
            stream_pairs[smallest.index] = IndexValue(smallest.index, None)
    return merged_list

print(merge_streams_better([iter([1, 4, 9]), iter([2, 6, 12]), iter([0, 4, 8])]))
print(merge_streams_better([iter([1, 4, 9])]))
print(merge_streams_better([iter([])]))