import operator import itertools class _AStopped(Exception): pass class _BStopped(Exception): pass def _raiser(excep): raise excep yield None 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 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): n = len(iterables) if n == 0: return iter(()) elif n == 1: return iter(iterables[0]) n2 = (n + 1) // 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)