Added new heapreplace(heap, item) function, to pop (and return) the

currently-smallest value, and add item, in one gulp.  See the second
N-Best algorithm in the test suite for a natural use.

2 files changed, 20 insertions, 3 deletions

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 278ba24..23f8be5 100644
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -14,6 +14,8 @@ heappush(heap, item) # pushes a new item on the heap
 item = heappop(heap) # pops the smallest item from the heap
 item = heap[0]       # smallest item on the heap without popping it
 heapify(x)           # transforms list into a heap, in-place, in linear time
+item = heapreplace(heap, item) # pops and returns smallest item, and adds
+                               # new item; the heap size is unchanged
 
 Our API differs from textbook heap algorithms as follows:
 
@@ -140,6 +142,22 @@ def heappop(heap):
         returnitem = lastelt
     return returnitem
 
+def heapreplace(heap, item):
+    """Pop and return the current smallest value, and add the new item.
+
+    This is more efficient than heappop() followed by heappush(), and can be
+    more appropriate when using a fixed-size heap.  Note that the value
+    returned may be larger than item!  That constrains reasonable uses of
+    this routine.
+    """
+
+    if heap:
+        returnitem = heap[0]
+        heap[0] = item
+        _siftup(heap, 0)
+        return returnitem
+    heap.pop()  # raise IndexError
+
 def heapify(x):
     """Transform list into a heap, in-place, in O(len(heap)) time."""
     n = len(x)
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 7f6d918..f0bb2e7 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -2,7 +2,7 @@
 
 from test.test_support import verify, vereq, verbose, TestFailed
 
-from heapq import heappush, heappop, heapify
+from heapq import heappush, heappop, heapify, heapreplace
 import random
 
 def check_invariant(heap):
@@ -68,8 +68,7 @@ def test_main():
     heapify(heap)
     for item in data[10:]:
         if item > heap[0]:  # this gets rarer the longer we run
-            heappop(heap)       # we know heap[0] isn't in best 10 anymore
-            heappush(heap, item)
+            heapreplace(heap, item)
     vereq(list(heapiter(heap)), data_sorted[-10:])
     # Make user happy
     if verbose: