1 files changed, 68 insertions, 30 deletions

diff --git a/Lib/heapq.py b/Lib/heapq.py
index dec15aec9c..d615239b94 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -127,8 +127,7 @@ From all times, sorting has always been a Great Art! :-)
 __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
            'nlargest', 'nsmallest', 'heappushpop']
 
-from itertools import islice, repeat, count, tee, chain
-import bisect
+from itertools import islice, count, tee, chain
 
 def heappush(heap, item):
     """Push item onto heap, maintaining the heap invariant."""
@@ -180,6 +179,19 @@ def heapify(x):
     for i in reversed(range(n//2)):
         _siftup(x, i)
 
+def _heappushpop_max(heap, item):
+    """Maxheap version of a heappush followed by a heappop."""
+    if heap and item < heap[0]:
+        item, heap[0] = heap[0], item
+        _siftup_max(heap, 0)
+    return item
+
+def _heapify_max(x):
+    """Transform list into a maxheap, in-place, in O(len(x)) time."""
+    n = len(x)
+    for i in reversed(range(n//2)):
+        _siftup_max(x, i)
+
 def nlargest(n, iterable):
     """Find the n largest elements in a dataset.
 
@@ -205,30 +217,16 @@ def nsmallest(n, iterable):
     """
     if n < 0:
         return []
-    if hasattr(iterable, '__len__') and n * 10 <= len(iterable):
-        # For smaller values of n, the bisect method is faster than a minheap.
-        # It is also memory efficient, consuming only n elements of space.
-        it = iter(iterable)
-        result = sorted(islice(it, 0, n))
-        if not result:
-            return result
-        insort = bisect.insort
-        pop = result.pop
-        los = result[-1]    # los --> Largest of the nsmallest
-        for elem in it:
-            if elem < los:
-                insort(result, elem)
-                pop()
-                los = result[-1]
+    it = iter(iterable)
+    result = list(islice(it, n))
+    if not result:
         return result
-    # An alternative approach manifests the whole iterable in memory but
-    # saves comparisons by heapifying all at once.  Also, saves time
-    # over bisect.insort() which has O(n) data movement time for every
-    # insertion.  Finding the n smallest of an m length iterable requires
-    #    O(m) + O(n log m) comparisons.
-    h = list(iterable)
-    heapify(h)
-    return list(map(heappop, repeat(h, min(n, len(h)))))
+    _heapify_max(result)
+    _heappushpop = _heappushpop_max
+    for elem in it:
+        _heappushpop(result, elem)
+    result.sort()
+    return result
 
 # 'heap' is a heap at all indices >= startpos, except possibly for pos.  pos
 # is the index of a leaf with a possibly out-of-order value.  Restore the
@@ -306,6 +304,42 @@ def _siftup(heap, pos):
     heap[pos] = newitem
     _siftdown(heap, startpos, pos)
 
+def _siftdown_max(heap, startpos, pos):
+    'Maxheap variant of _siftdown'
+    newitem = heap[pos]
+    # Follow the path to the root, moving parents down until finding a place
+    # newitem fits.
+    while pos > startpos:
+        parentpos = (pos - 1) >> 1
+        parent = heap[parentpos]
+        if parent < newitem:
+            heap[pos] = parent
+            pos = parentpos
+            continue
+        break
+    heap[pos] = newitem
+
+def _siftup_max(heap, pos):
+    'Maxheap variant of _siftup'
+    endpos = len(heap)
+    startpos = pos
+    newitem = heap[pos]
+    # Bubble up the larger child until hitting a leaf.
+    childpos = 2*pos + 1    # leftmost child position
+    while childpos < endpos:
+        # Set childpos to index of larger child.
+        rightpos = childpos + 1
+        if rightpos < endpos and not heap[rightpos] < heap[childpos]:
+            childpos = rightpos
+        # Move the larger child up.
+        heap[pos] = heap[childpos]
+        pos = childpos
+        childpos = 2*pos + 1
+    # The leaf at pos is empty now.  Put newitem there, and bubble it up
+    # to its final resting place (by sifting its parents down).
+    heap[pos] = newitem
+    _siftdown_max(heap, startpos, pos)
+
 # If available, use C implementation
 try:
     from _heapq import *
@@ -324,6 +358,7 @@ def merge(*iterables):
 
     '''
     _heappop, _heapreplace, _StopIteration = heappop, heapreplace, StopIteration
+    _len = len
 
     h = []
     h_append = h.append
@@ -335,17 +370,20 @@ def merge(*iterables):
             pass
     heapify(h)
 
-    while 1:
+    while _len(h) > 1:
         try:
-            while 1:
-                v, itnum, next = s = h[0]   # raises IndexError when h is empty
+            while True:
+                v, itnum, next = s = h[0]
                 yield v
                 s[0] = next()               # raises StopIteration when exhausted
                 _heapreplace(h, s)          # restore heap condition
         except _StopIteration:
             _heappop(h)                     # remove empty iterator
-        except IndexError:
-            return
+    if h:
+        # fast case when only a single iterator remains
+        v, itnum, next = h[0]
+        yield v
+        yield from next.__self__
 
 # Extend the implementations of nsmallest and nlargest to use a key= argument
 _nsmallest = nsmallest