Minor fiddling, including a simple class to implement a heap iterator

tim-one · tim-one · commit aa7d24319ebf · 2002-08-03T02:11:26.000Z
in the test file.  I have docs for heapq.heapify ready to check in, but
Jack appears to have left behind a stale lock in the Doc/lib directory.
diff --git a/Lib/heapq.py b/Lib/heapq.py
@@ -13,7 +13,7 @@
 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(heap)        # transform list into a heap, in-place, in linear time
+heapify(x)           # transforms list into a heap, in-place, in linear time
 
 Our API differs from textbook heap algorithms as follows:
 
@@ -175,16 +175,16 @@ def heappop(heap):
         returnitem = lastelt
     return returnitem
 
-def heapify(heap):
-    """Transform list heap into a heap, in-place, in O(len(heap)) time."""
-    n = len(heap)
+def heapify(x):
+    """Transform list into a heap, in-place, in O(len(heap)) time."""
+    n = len(x)
     # Transform bottom-up.  The largest index there's any point to looking at
     # is the largest with a child index in-range, so must have 2*i + 1 < n,
     # or i < (n-1)/2.  If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so
     # j-1 is the largest, which is n//2 - 1.  If n is odd = 2*j+1, this is
     # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1.
     for i in xrange(n//2 - 1, -1, -1):
-        _siftdown(heap, i)
+        _siftdown(x, i)
 
 if __name__ == "__main__":
     # Simple sanity test
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
@@ -12,6 +12,20 @@ def check_invariant(heap):
             parentpos = (pos-1) >> 1
             verify(heap[parentpos] <= item)
 
+# An iterator returning a heap's elements, smallest-first.
+class heapiter(object):
+    def __init__(self, heap):
+        self.heap = heap
+
+    def next(self):
+        try:
+            return heappop(self.heap)
+        except IndexError:
+            raise StopIteration
+
+    def __iter__(self):
+        return self
+
 def test_main():
     # 1) Push 100 random numbers and pop them off, verifying all's OK.
     heap = []
@@ -47,17 +61,16 @@ def test_main():
         check_invariant(heap)
     # 5) Less-naive "N-best" algorithm, much faster (if len(data) is big
     #    enough <wink>) than sorting all of data.  However, if we had a max
-    #    heap instead of a min heap, it would go much faster still via
+    #    heap instead of a min heap, it could go faster still via
     #    heapify'ing all of data (linear time), then doing 10 heappops
     #    (10 log-time steps).
     heap = data[:10]
     heapify(heap)
     for item in data[10:]:
-        if item > heap[0]:  # this gets rarer and rarer the longer we run
+        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)
-            heappop(heap)
-    heap.sort()
-    vereq(heap, data_sorted[-10:])
+    vereq(list(heapiter(heap)), data_sorted[-10:])
     # Make user happy
     if verbose:
         print "All OK"
