"merge() implemented with a binary-tree of iterators"


class Node:
    # Invariants:
    # * top node: parent is None
    # * leaf node: left and right are None, and source is an Callable
    # * non-leaf nodes:  left and right are not None, and source is a Node
    __slots__ = "value", "source", "parent", "left", "right"

    def __init__(self, value, source, left=None, right=None):
        self.value = value
        self.source = source
        self.left = left
        self.right = right
        self.parent = None

def merge(*iterables):
    # build leaf nodes
    sentinel = object()
    nodes = []
    for iterable in iterables:
        it = iter(iterable)
        value = next(it, sentinel)
        if value is sentinel:
            continue
        nodes += [Node(value, it.__next__)]
    if not nodes:
        return

    # Build tree of non-leaf nodes
    while len(nodes) > 1:
        newnodes = []
        for c1, c2 in zip(nodes[::2], nodes[1::2]):
            selection = c1 if c1.value < c2.value else c2
            node = Node(selection.value, selection, c1, c2)
            c1.parent = c2.parent = node
            newnodes += [node]
        if len(nodes) & 1:
            newnodes += nodes[-1:]
        nodes = newnodes
    node = nodes.pop()

    while True:
        # yield top value
        yield node.value

        # follow down the chain of sources to an iterator
        while node.left is not None:
            node = node.source

        # load the next value from the iterator
        try:
            node.value = node.source()
        except StopIteration:
            if node.parent is None:
                return
            # when empty, move sibling node up a level
            parent = node.parent
            sibling = parent.left if parent.right is node else parent.right
            parent.value = sibling.value
            parent.source = sibling.source
            parent.left = sibling.left
            parent.right = sibling.right
            node = parent

        # follow up the chain of parent branch nodes
        while node.parent is not None:
            node = node.parent
            c1 = node.left
            c2 = node.right
            selection = c1 if c1.value < c2.value else c2
            node.value = selection.value
            node.source = selection


if __name__ == "__main__":

    arr0 = [10, 20, 30, 40, 50, 60]
    arr1 = [15, 17, 35, 38]
    arr2 = [12, 32]
    arr3 = [6, 22, 25, 61]
    arr4 = []
    expected = sorted(arr0 + arr1 + arr2 + arr3 + arr4)
    actual = list(merge(arr0, arr1, arr2, arr3, arr4))
    print(expected)
    print(actual)
    assert expected == actual