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     # Make user happy
  74     if verbose:
  75         print "All OK"
  76
  77 if __name__ == "__main__":
  78     test_main()