diff options
Diffstat (limited to 'Lib/test/test_heapq.py')
| -rw-r--r-- | Lib/test/test_heapq.py | 176 |
1 files changed, 90 insertions, 86 deletions
diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index f04ea94..944b17d 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -1,101 +1,105 @@ """Unittests for heapq.""" -from test.test_support import verify, vereq, verbose, TestFailed - from heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest import random +import unittest +from test import test_support -def check_invariant(heap): - # Check the heap invariant. - for pos, item in enumerate(heap): - if pos: # pos 0 has no parent - 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 heapiter(heap): + # An iterator returning a heap's elements, smallest-first. + try: + while 1: + yield heappop(heap) + except IndexError: + pass - def next(self): - try: - return heappop(self.heap) - except IndexError: - raise StopIteration +class TestHeap(unittest.TestCase): - def __iter__(self): - return self + def test_push_pop(self): + # 1) Push 256 random numbers and pop them off, verifying all's OK. + heap = [] + data = [] + self.check_invariant(heap) + for i in range(256): + item = random.random() + data.append(item) + heappush(heap, item) + self.check_invariant(heap) + results = [] + while heap: + item = heappop(heap) + self.check_invariant(heap) + results.append(item) + data_sorted = data[:] + data_sorted.sort() + self.assertEqual(data_sorted, results) + # 2) Check that the invariant holds for a sorted array + self.check_invariant(results) -def test_main(): - # 1) Push 100 random numbers and pop them off, verifying all's OK. - heap = [] - data = [] - check_invariant(heap) - for i in range(256): - item = random.random() - data.append(item) - heappush(heap, item) - check_invariant(heap) - results = [] - while heap: - item = heappop(heap) - check_invariant(heap) - results.append(item) - data_sorted = data[:] - data_sorted.sort() - vereq(data_sorted, results) - # 2) Check that the invariant holds for a sorted array - check_invariant(results) - # 3) Naive "N-best" algorithm - heap = [] - for item in data: - heappush(heap, item) - if len(heap) > 10: - heappop(heap) - heap.sort() - vereq(heap, data_sorted[-10:]) - # 4) Test heapify. - for size in range(30): - heap = [random.random() for dummy in range(size)] - heapify(heap) - 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 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 the longer we run - heapreplace(heap, item) - vereq(list(heapiter(heap)), data_sorted[-10:]) - # 6) Exercise everything with repeated heapsort checks - for trial in xrange(100): - size = random.randrange(50) - data = [random.randrange(25) for i in range(size)] - if trial & 1: # Half of the time, use heapify - heap = data[:] + def check_invariant(self, heap): + # Check the heap invariant. + for pos, item in enumerate(heap): + if pos: # pos 0 has no parent + parentpos = (pos-1) >> 1 + self.assert_(heap[parentpos] <= item) + + def test_heapify(self): + for size in range(30): + heap = [random.random() for dummy in range(size)] heapify(heap) - else: # The rest of the time, use heappush - heap = [] - for item in data: - heappush(heap,item) - data.sort() - sorted = [heappop(heap) for i in range(size)] - vereq(data, sorted) + self.check_invariant(heap) + + def test_naive_nbest(self): + data = [random.randrange(2000) for i in range(1000)] + heap = [] + for item in data: + heappush(heap, item) + if len(heap) > 10: + heappop(heap) + heap.sort() + self.assertEqual(heap, sorted(data)[-10:]) - # 7) Check nlargest() and nsmallest() - data = [random.randrange(2000) for i in range(1000)] - copy = data[:] - copy.sort(reverse=True) - vereq(nlargest(data, 400), copy[:400]) - copy.sort() - vereq(nsmallest(data, 400), copy[:400]) + def test_nbest(self): + # 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 could go faster still via + # heapify'ing all of data (linear time), then doing 10 heappops + # (10 log-time steps). + data = [random.randrange(2000) for i in range(1000)] + heap = data[:10] + heapify(heap) + for item in data[10:]: + if item > heap[0]: # this gets rarer the longer we run + heapreplace(heap, item) + self.assertEqual(list(heapiter(heap)), sorted(data)[-10:]) - # Make user happy - if verbose: - print "All OK" + def test_heapsort(self): + # Exercise everything with repeated heapsort checks + for trial in xrange(100): + size = random.randrange(50) + data = [random.randrange(25) for i in range(size)] + if trial & 1: # Half of the time, use heapify + heap = data[:] + heapify(heap) + else: # The rest of the time, use heappush + heap = [] + for item in data: + heappush(heap, item) + heap_sorted = [heappop(heap) for i in range(size)] + self.assertEqual(heap_sorted, sorted(data)) + + def test_nsmallest(self): + data = [random.randrange(2000) for i in range(1000)] + self.assertEqual(nsmallest(data, 400), sorted(data)[:400]) + + def test_largest(self): + data = [random.randrange(2000) for i in range(1000)] + self.assertEqual(nlargest(data, 400), sorted(data, reverse=True)[:400]) + +def test_main(): + test_support.run_unittest(TestHeap) if __name__ == "__main__": test_main() + |