Merges iterators and presumes ordering

flatten_soted_iterators.py

import heapq


class HeapElem:
    def __init__(self, val: int, iter):
        self.iter = iter
        self.val = val

    def __lt__(self, other):
        return self.val < other.val


class Solution:
    def __init__(self, *argv):
        self.hp = []
        self.iterators = list(argv)

        for iter in self.iterators:
            self.push_next(iter)

    def push_next(self, arg):
        try:
            heapq.heappush(self.hp, HeapElem(val=next(arg), iter=arg))
        except StopIteration:
            pass

    def __iter__(self):
        return self

    def __next__(self):
        if not self.hp:
            raise StopIteration

        min_elem = heapq.heappop(self.hp)

        if not self.iterators:
            return min_elem.val

        self.push_next(min_elem.iter)
        return min_elem.val


if __name__ == '__main__':
    input = [[5, 8, 9], [10, 11, 12], [13, 14, 15], [16, 17, 18], [0, 1, 2, 3, 6, 7]]
    iter = [iter(a) for a in input]

    for elem in Solution(*iter):
        print(elem)