Add merge() function to heapq.

1 files changed, 40 insertions, 2 deletions

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 753c3b7..b56d0f9 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -126,8 +126,8 @@ Believe me, real good tape sorts were quite spectacular to watch!
 From all times, sorting has always been a Great Art! :-)
 """
 
-__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'nlargest',
-           'nsmallest']
+__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
+           'nlargest', 'nsmallest']
 
 from itertools import islice, repeat, count, imap, izip, tee
 from operator import itemgetter, neg
@@ -308,6 +308,41 @@ try:
 except ImportError:
     pass
 
+def merge(*iterables):
+    '''Merge multiple sorted inputs into a single sorted output.
+
+    Similar to sorted(itertools.chain(*iterables)) but returns an iterable,
+    does not pull the data into memory all at once, and reduces the number
+    of comparisons by assuming that each of the input streams is already sorted.
+
+    >>> 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]
+
+    '''
+    _heappop, siftup, _StopIteration = heappop, _siftup, StopIteration
+
+    h = []
+    h_append = h.append
+    for it in map(iter, iterables):
+        try:
+            next = it.next
+            h_append([next(), next])
+        except _StopIteration:
+            pass
+    heapify(h)
+
+    while 1:
+        try:
+            while 1:
+                v, next = s = h[0]      # raises IndexError when h is empty
+                yield v
+                s[0] = next()           # raises StopIteration when exhausted
+                siftup(h, 0)            # restore heap condition
+        except _StopIteration:
+            _heappop(h)                  # remove empty iterator
+        except IndexError:
+            return
+
 # Extend the implementations of nsmallest and nlargest to use a key= argument
 _nsmallest = nsmallest
 def nsmallest(n, iterable, key=None):
@@ -341,3 +376,6 @@ if __name__ == "__main__":
     while heap:
         sort.append(heappop(heap))
     print sort
+
+    import doctest
+    doctest.testmod()