Adding the heap queue algorithm, per discussion in python-dev last

week.

1 files changed, 48 insertions, 0 deletions

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
new file mode 100644
index 0000000000..43723f33f6
--- /dev/null
+++ b/Lib/test/test_heapq.py
@@ -0,0 +1,48 @@
+"""Unittests for heapq."""
+
+from test.test_support import verify, vereq, verbose, TestFailed
+
+from heapq import heappush, heappop
+import random
+
+def check_invariant(heap):
+    # Check the heap invariant.
+    for pos, item in enumerate(heap):
+        parentpos = (pos+1)/2 - 1
+        if parentpos >= 0:
+            verify(heap[parentpos] <= item)
+
+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:])
+    # Make user happy
+    if verbose:
+        print "All OK"
+
+if __name__ == "__main__":
+    test_main()