From aa7d24319ebf62d38463e798b99be1afad314db6 Mon Sep 17 00:00:00 2001 From: Tim Peters Date: Sat, 3 Aug 2002 02:11:26 +0000 Subject: Minor fiddling, including a simple class to implement a heap iterator 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. --- Lib/test/test_heapq.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) (limited to 'Lib/test') diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 1330f1249c..7f6d918b51 100644 --- 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 ) 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" -- cgit v1.2.1