diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/heapq.py | 18 | ||||
| -rw-r--r-- | Lib/test/test_heapq.py | 5 |
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: |