class Node: __slots__ = "value", "source_index", "parent", "left", "right" def __init__(self, value, source_index): self.value = value self.source_index = source_index self.parent = self.left = self.right = None def treeify(pairs): n = len(pairs) if n == 1: return pairs[0] n = (n + 1) // 2 left_root, left_champion = treeify(pairs[:n]) right_root, right_champion = treeify(pairs[n:]) if right_champion.value < left_champion.value: loser, winner = left_champion, right_champion else: winner, loser = left_champion, right_champion left_root.parent = right_root.parent = loser loser.left = left_root loser.right = right_root return loser, winner def merge(*iterables): sentinel = object() _next = next pairs = [] leaves = [] for iterable in iterables: it = iter(iterable) value = _next(it, sentinel) if value is not sentinel: leaf = Node(it, "it"+str(len(leaves))) #DEBUG pairs.append((leaf, Node(value, len(leaves)))) leaves.append(leaf) if not pairs: return node, champion = treeify(pairs) in_hand_index = champion.source_index in_hand_value = champion.value while True: yield in_hand_value leaf = leaves[in_hand_index] in_hand_value = _next(leaf.value, sentinel) if in_hand_value is sentinel: parent = leaf.parent if parent is None: return in_hand_value = parent.value in_hand_index = parent.source_index left, right = parent.left, parent.right sibling = left if leaf is right else right grandparent = parent.parent sibling.parent = grandparent if sibling.left is None and grandparent is None: # merge(leaf, sibling) --> sibling yield in_hand_value yield from sibling.value return if grandparent is None: pass elif grandparent.left is parent: grandparent.left = sibling else: grandparent.right = sibling node = sibling else: node = leaf while (node := node.parent) is not None: node_index = node.source_index node_value = node.value if in_hand_index < node_index: if not node_value < in_hand_value: continue else: if in_hand_value < node_value: continue # found a better one, so pick it up; # swap node with in_hand node.source_index = in_hand_index in_hand_index = node_index node.value = in_hand_value in_hand_value = node_value 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 import random for i in range(100_000): its = [sorted(random.randint(-5,5) for _ in range(random.randrange(5))) for _ in range(random.randrange(10))] expected = sorted(sum(its, [])) actual = list(merge(*its)) if expected != actual: print(expected) print(actual) else: print(i) # def merge(*iterables): # iters = [iter(x) for x in iterables] # n = len(iters) # if n == 0: # return # if n == 1: # yield from iters[0] # return # min_sentinel = object() # sentinel = object() # tree = [min_sentinel] * (n-1) # _StopIteration = StopIteration # _next = next # for iter_index, it in enumerate(iters): # try: # held_node = (_next(it), iter_index) # except _StopIteration: # held_node = sentinel # tree_index = (iter_index + n - 2) // 2 # while True: # traverse = tree[tree_index] # if traverse is min_sentinel: # tree[tree_index] = held_node # if held_node is sentinel \ # or traverse is min_sentinel \ # or (held_node is not min_sentinel # and traverse is not sentinel # and traverse < held_node): # tree[tree_index] = held_node # held_node = traverse # if tree_index == 0: # break # tree_index = (tree_index - 1) // 2 # assert held_node is min_sentinel or iter_index == n - 1 # while True: # if held_node is sentinel: # return # champion, iter_index = held_node # yield champion # try: # held_node = (_next(iters[iter_index]), iter_index) # except _StopIteration: # held_node = sentinel # tree_index = (iter_index + n - 2) // 2 # while True: # traverse = tree[tree_index] # if held_node is sentinel \ # or (traverse is not sentinel # and traverse < held_node): # tree[tree_index] = held_node # held_node = traverse # if tree_index == 0: # break # tree_index = (tree_index - 1) // 2