home *** CD-ROM | disk | FTP | other *** search
/ PC World 2005 June / PCWorld_2005-06_cd.bin / software / vyzkuste / firewally / firewally.exe / framework-2.3.exe / test_heapq.py < prev    next >
Text File  |  2003-12-30  |  3KB  |  93 lines
  1. """Unittests for heapq."""
  2.  
  3. from test.test_support import verify, vereq, verbose, TestFailed
  4.  
  5. from heapq import heappush, heappop, heapify, heapreplace
  6. import random
  7.  
  8. def check_invariant(heap):
  9.     # Check the heap invariant.
  10.     for pos, item in enumerate(heap):
  11.         if pos: # pos 0 has no parent
  12.             parentpos = (pos-1) >> 1
  13.             verify(heap[parentpos] <= item)
  14.  
  15. # An iterator returning a heap's elements, smallest-first.
  16. class heapiter(object):
  17.     def __init__(self, heap):
  18.         self.heap = heap
  19.  
  20.     def next(self):
  21.         try:
  22.             return heappop(self.heap)
  23.         except IndexError:
  24.             raise StopIteration
  25.  
  26.     def __iter__(self):
  27.         return self
  28.  
  29. def test_main():
  30.     # 1) Push 100 random numbers and pop them off, verifying all's OK.
  31.     heap = []
  32.     data = []
  33.     check_invariant(heap)
  34.     for i in range(256):
  35.         item = random.random()
  36.         data.append(item)
  37.         heappush(heap, item)
  38.         check_invariant(heap)
  39.     results = []
  40.     while heap:
  41.         item = heappop(heap)
  42.         check_invariant(heap)
  43.         results.append(item)
  44.     data_sorted = data[:]
  45.     data_sorted.sort()
  46.     vereq(data_sorted, results)
  47.     # 2) Check that the invariant holds for a sorted array
  48.     check_invariant(results)
  49.     # 3) Naive "N-best" algorithm
  50.     heap = []
  51.     for item in data:
  52.         heappush(heap, item)
  53.         if len(heap) > 10:
  54.             heappop(heap)
  55.     heap.sort()
  56.     vereq(heap, data_sorted[-10:])
  57.     # 4) Test heapify.
  58.     for size in range(30):
  59.         heap = [random.random() for dummy in range(size)]
  60.         heapify(heap)
  61.         check_invariant(heap)
  62.     # 5) Less-naive "N-best" algorithm, much faster (if len(data) is big
  63.     #    enough <wink>) than sorting all of data.  However, if we had a max
  64.     #    heap instead of a min heap, it could go faster still via
  65.     #    heapify'ing all of data (linear time), then doing 10 heappops
  66.     #    (10 log-time steps).
  67.     heap = data[:10]
  68.     heapify(heap)
  69.     for item in data[10:]:
  70.         if item > heap[0]:  # this gets rarer the longer we run
  71.             heapreplace(heap, item)
  72.     vereq(list(heapiter(heap)), data_sorted[-10:])
  73.     # 6)  Exercise everything with repeated heapsort checks
  74.     for trial in xrange(100):
  75.         size = random.randrange(50)
  76.         data = [random.randrange(25) for i in range(size)]
  77.         if trial & 1:     # Half of the time, use heapify
  78.             heap = data[:]
  79.             heapify(heap)
  80.         else:             # The rest of the time, use heappush
  81.             heap = []
  82.             for item in data:
  83.                 heappush(heap,item)
  84.         data.sort()
  85.         sorted = [heappop(heap) for i in range(size)]
  86.         vereq(data, sorted)
  87.     # Make user happy
  88.     if verbose:
  89.         print "All OK"
  90.  
  91. if __name__ == "__main__":
  92.     test_main()
  93.