Lib/test/test_heapq.py

   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()