from heapq import heappop, heapreplace, heapify def imerge(iterables, key=None): '''Merge multiple (potentially infinite) sorted inputs into a single sorted output. Similar to heapq.merge and sorted(itertools.chain(*iterables)). Like heapq.merge, returns a generator, does not pull the data into memory all at once, and assumes that each of the input iterables is already sorted (smallest to largest). Unlike heapq.merge, accepts an infinite number of input iterables, but requires all of them to come in ascending order (that is, their starting point must be in ascending order each). >>> list(imerge([[1,3,5,7], [2,4,8], [5,10,15,20], [], [25]])) [1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25] ''' _heappop, _heapreplace, _heapify, _StopIteration = heappop, heapreplace, heapify, StopIteration h = [] h_append = h.append iterables = iter(iterables) more_iterables = True while len(h)<2: try: # raises StopIteration when no more iterables next_item = iter(next(iterables)).next except _StopIteration: more_iterables = False break try: v = next_item() except _StopIteration: # ignore empty iterables continue if key is not None: highest = key(v) else: highest = v h_append([highest, v, next_item]) if len(h) >= 2: # the heap invariant should hold, if input iterables come already sorted # XXX: make this a ValueError assert h[1][0] >= h[0][0] elif len(h) == 1: # a single iterable, just send it assert not more_iterables _, v, next_item = h[0] yield v try: while True: yield next_item() except _StopIteration: return else: # empty return cur = highest while h: _, v, next_item = s = h[0] yield v try: v = s[1] = next_item() # raises StopIteration when no more items except _StopIteration: _heappop(h) # remove empty iterator else: if key is not None: cur = s[0] = key(v) else: cur = s[0] = v _heapreplace(h, s) # restore heap condition # 'highest' is the highest known item in the heap. # Any time we advance an iterable and get an item ('cur') # greater than 'highest', we must bring more enough iterables # into play to ensure no items are missed. if more_iterables and (cur >= highest or len(h) < 2): while cur >= highest or len(h)<2: try: # raises StopIteration when no more iterables next_item = iter(next(iterables)).next except _StopIteration: more_iterables = False break try: v = next_item() except _StopIteration: # ignore empty iterables continue if key is not None: highest = key(v) else: highest = v h_append([highest, v, next_item]) _heapify(h)