def merge(*iterables, key=None, reverse=False): '''Merge multiple sorted inputs into a single sorted output. Similar to sorted(itertools.chain(*iterables)) but returns a generator, does not pull the data into memory all at once, and assumes that each of the input streams is already sorted (smallest to largest). >>> list(merge([1,3,5,7], [0,2,4,8], [5,10,15,20], [], [25])) [0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25] If *key* is not None, applies a key function to each element to determine its sort order. >>> list(merge(['dog', 'horse'], ['cat', 'fish', 'kangaroo'], key=len)) ['dog', 'cat', 'fish', 'horse', 'kangaroo'] ''' # Invariants for the binary tree structure used here: # - All leaf nodes have leaf[LEFT] = None and leaf[RIGHT] = the # __next__ method of an iterator. # - For non-leaf nodes, node[LEFT][PARENT] is node and # node[RIGHT][PARENT] is node. # - The root's PARENT is None, and all other PARENTs are nodes. # - Each node's KEY and VALUE are identical to the KEY and VALUE # of one of its descendent leaves. A node's LEAF points to # which descendent leaf the key and value come from. # - A parent's key/value will be the same as the higher-priority # key/value among that of its two children # A node will just be a list with these semantic aliases for # indices into it. PARENT, LEFT, RIGHT, LEAF, VALUE = range(5) KEY = VALUE if key is None else 5 sentinel = object() # Produce leaf nodes, storing a value/key its source. # This eliminates empty iterators. nodes = [] for it in map(iter, iterables): value = next(it, sentinel) if value is not sentinel: leaf = [None, None, it.__next__, None, value] if key is not None: leaf.append(key(value)) leaf[LEAF] = leaf nodes.append(leaf) # fast early-outs if not nodes: return if len(nodes) == 1: node = nodes[0] yield node[VALUE] yield from node[RIGHT].__self__ return # Add a shared parent for adjacent nodes, then add a shared parent # for adjacent parents, etc., until there is one common root. while (n := len(nodes)) > 1: # split odd numbers in favor of the right (leaving the # unpaired nodes to the left). This is better because, e.g.: # - merge([1], merge([2], [1])) requires 2 comparisons, but # - merge(merge([1], [2]), [1]) requires 3 comparisons due # to stability considerations. new_nodes = nodes[:n & 1] for i in range(n & 1, n - 1, 2): left, right = nodes[i:i + 2] if reverse: winner = right if left[KEY] < right[KEY] else left else: winner = right if right[KEY] < left[KEY] else left node = [None, left, right] # the winner's LEAF, VALUE, and KEY are copied into the # parent node node[LEAF:] = winner[LEAF:] left[PARENT] = right[PARENT] = node new_nodes.append(node) nodes = new_nodes (root,) = nodes del nodes # main loop while True: yield root[VALUE] # jump to the leaf that produced the value we just yielded. node = root[LEAF] # replace the leaf's value using the iterator try: node[VALUE] = node[RIGHT]() except StopIteration: # Upon exhuastion, Promote the exhausted leaf's sibling. # The sibling must exist because of the early-outs. parent = node[PARENT] left, right = parent[LEFT], parent[RIGHT] sibling = left if right is node else right # copy everything but the PARENT parent[LEFT:] = sibling[LEFT:] if sibling[LEFT] is not None: # If sibling has children, give custody to parent. sibling[LEFT][PARENT] = sibling[RIGHT][PARENT] = parent parent[LEAF] = sibling[LEAF] else: # If sibling is a leaf, then parent is now a leaf. parent[LEAF] = parent if parent is root: # a leaf has been promoted to root. # There's only one node, so do a fast out. yield parent[VALUE] yield from parent[RIGHT].__self__ return # break some GC cycles del sibling[:], sibling node = parent else: # If a new value was successfully supplied, compute its key. if key is not None: node[KEY] = key(node[VALUE]) # "replay the games" leading from the leaf to the root, i.e., # restore the invariant that each parent's value coicides with # its higher-priority child's value. if reverse: while node is not root: node = node[PARENT] left, right = node[LEFT], node[RIGHT] winner = right if left[KEY] < right[KEY] else left # advance the winner's LEAF, VALUE, and KEY node[LEAF:] = winner[LEAF:] else: while node is not root: node = node[PARENT] left, right = node[LEFT], node[RIGHT] winner = right if right[KEY] < left[KEY] else left # advance the winner's LEAF, VALUE, and KEY node[LEAF:] = winner[LEAF:]