import operator import itertools def _merge_two(a_iter, b_iter, key, reverse): a_iter = iter(a_iter) b_iter = iter(b_iter) try: a = next(a_iter) except StopIteration: yield from b_iter return class AStopped(Exception): pass class BStopped(Exception): pass def raiser(excep): raise excep yield None next_a = itertools.chain(a_iter, raiser(AStopped)).__next__ next_b = itertools.chain(b_iter, raiser(BStopped)).__next__ try: b = next_b() if key is None: if not reverse: while True: if b < a: yield b b = next_b() else: yield a a = next_a() else: while True: if a < b: yield b b = next_b() else: yield a a = next_a() else: key_a = key(a) key_b = key(b) if not reverse: while True: if key_b < key_a: yield b b = next_b() key_b = key(b) else: yield a a = next_a() key_a = key(a) else: while True: if key_a < key_b: yield b b = next_b() key_b = key(b) else: yield a a = next_a() key_a = key(a) except AStopped: yield b yield from b_iter except BStopped: yield a yield from a_iter _merge_recurse = object() _merge_get_value = operator.itemgetter(0) _merge_get_key = operator.itemgetter(1) 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'] ''' n = len(iterables) if n == 0: return iter(()) elif n == 1: return iter(iterables[0]) elif n == 2: if key is _merge_recurse: key = _merge_get_key return _merge_two(*iterables, key=key, reverse=reverse) n2 = n // 2 if key is None: left = merge(*iterables[:n2], key=None, reverse=reverse) right = merge(*iterables[n2:], key=None, reverse=reverse) return _merge_two(left, right, key=None, reverse=reverse) elif key is _merge_recurse: # recursive call, keys already computed left = merge(*iterables[:n2], key=_merge_recurse, reverse=reverse) right = merge(*iterables[n2:], key=_merge_recurse, reverse=reverse) return _merge_two(left, right, key=_merge_get_key, reverse=reverse) else: # Top-level call: zip iterators up with their keys once and for all. # Then key computation is just an itemgetter. key_iterables = [zip(it, map(key, it)) for it in iterables] left = merge(*key_iterables[:n2], key=_merge_recurse, reverse=reverse) right = merge(*key_iterables[n2:], key=_merge_recurse, reverse=reverse) result = _merge_two(left, right, key=_merge_get_key, reverse=reverse) return map(_merge_get_value, result)