chain_ordered.py

from typing import Iterable, Iterator, Any


def chain_ordered[T](
        *seqs: Iterable[T],
) -> Iterable[T]:
     """
    Chain ordered sequences into ordered sequence. Complexity of one
    iteration is O(n), where n is a number of sequences.

    >>> import itertools
    >>> seqs = [
    ...     [1, 2, 3],
    ...     [-5, 6, 12],
    ... ]
    >>> assert list(chain_ordered(*seqs)) == sorted(itertools.chain(*seqs))

    """

    competing_elements: list[T] = []
    iterators: list[Iterator[T]] = [iter(i) for i in seqs]
    sentinel = object()
    common_next: Any = sentinel

    for i in iterators.copy():
        try:
            next_item = next(i)
        except StopIteration:
            iterators.remove(i)
        else:
            competing_elements.append(next_item)
            if common_next is sentinel or common_next > next_item:
                common_next = next_item

    if common_next is sentinel:
        return

    while True:
        yield common_next

        common_next_first_index = competing_elements.index(common_next)

        try:
            competing_elements[common_next_first_index] = \
                next(iterators[common_next_first_index])
        except StopIteration:
            iterators.pop(common_next_first_index)
            competing_elements.pop(common_next_first_index)

        if not iterators:
            break

        common_next = min(competing_elements)