Lib/test/test_heapq.py

"""Unittests for heapq."""

from heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest
import random
import unittest
from test import test_support
import sys


def heapiter(heap):
    # An iterator returning a heap's elements, smallest-first.
    try:
        while 1:
            yield heappop(heap)
    except IndexError:
        pass

class TestHeap(unittest.TestCase):

    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 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)
            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:])

    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:])

    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)]
        for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
            self.assertEqual(nsmallest(n, data), sorted(data)[:n])

    def test_largest(self):
        data = [random.randrange(2000) for i in range(1000)]
        for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
            self.assertEqual(nlargest(n, data), sorted(data, reverse=True)[:n])

def test_main(verbose=None):
    test_classes = [TestHeap]
    test_support.run_unittest(*test_classes)

    # verify reference counting
    if verbose and hasattr(sys, "gettotalrefcount"):
        import gc
        counts = [None] * 5
        for i in xrange(len(counts)):
            test_support.run_unittest(*test_classes)
            gc.collect()
            counts[i] = sys.gettotalrefcount()
        print counts

if __name__ == "__main__":
    test_main(verbose=True)