def test_getitem(self): l = range(10, -1, -1) h = Heap() caught_exception = False try: print h[1] except HeapEmptyError: caught_exception = True assert caught_exception == True for i in xrange(len(l)): h.add(l[i]) caught_exception = False try: print h[-2] except IndexError: caught_exception = True assert caught_exception == True assert len(h) == 11 caught_exception = False try: print h[11] except IndexError: caught_exception = True assert caught_exception == True x = [0, 1, 5, 4, 2, 9, 6, 10, 7, 8, 3] for i in xrange(10): assert h[i] == x[i]
def test_custom_comparator(self): # Max heap using custom comparator h = Heap(lambda a,b: cmp(b, a)) for i in range(10): h.add(i) for i in range(9, -1, -1): assert h.remove() == i
def test_add(self): l = [5,4,3,2,1,6,7,8,9,0] intermediate_heaps = [ [5], [4, 5], [3, 5, 4], [2, 3, 4, 5], [1, 2, 4, 5, 3], [1, 2, 4, 5, 3, 6], [1, 2, 4, 5, 3, 6, 7], [1, 2, 4, 5, 3, 6, 7, 8], [1, 2, 4, 5, 3, 6, 7, 8, 9], [0, 1, 4, 5, 2, 6, 7, 8, 9, 3] ] h = Heap() for i in xrange(len(l)): h.add(l[i]) self.assertEqual (h.items, intermediate_heaps[i])
def test_remove(self): l = [5,4,3,2,1,6,7,8,9,0] h = Heap() with self.assertRaises(HeapEmptyError): h.peek() caught_exception = False try: assert h.remove() == None except HeapEmptyError as e: assert str(e) == "HeapEmptyError: 'remove(): Heap is empty'" caught_exception = True assert caught_exception == True for i in xrange(len(l)): h.add(l[i]) assert len(h) == 10 for i in xrange(10): item = h.remove() self.assertEqual(len(h), 10-i-1) self.assertEqual(item, i)
def test_custom_items(self):
class Record:
def __init__(self, a, b, c):
self.a = a
self.b = b
self.c = c
# by default, using 'b' as key to compare
def __cmp__(self, other):
return cmp(self.b, other.b)
def __str__(self):
return "(%s, %s, %s)" %(self.a, self.b, self.c)
h = Heap()
h.add(Record("record1", 1, 100))
h.add(Record("record4", 5, 125))
h.add(Record("record3", 2, 50))
h.add(Record("record2", 3, 25))
h.add(Record("record5", 4, 5))
sorted_records_bs = []
h.sorted(lambda x, l: l.append(x.b), sorted_records_bs)
assert sorted_records_bs == range(1,6)
h2 = Heap(lambda r1, r2: cmp(r1.a, r2.a))
h2.add(Record("record1", 1, 100))
h2.add(Record("record4", 5, 125))
h2.add(Record("record3", 2, 50))
h2.add(Record("record2", 3, 25))
h2.add(Record("record5", 4, 5))
sorted_records_as = []
h2.sorted(lambda x, l: l.append(x.a), sorted_records_as)
assert sorted_records_as == ["record"+str(i) for i in xrange(1,6)]
h3 = Heap(lambda r1, r2: cmp(r1.c, r2.c))
h3.add(Record("record1", 1, 100))
h3.add(Record("record4", 5, 125))
h3.add(Record("record3", 2, 50))
h3.add(Record("record2", 3, 25))
h3.add(Record("record5", 4, 5))
sorted_records_cs = []
h3.sorted(lambda x, l: l.append(x.c), sorted_records_cs)
assert sorted_records_cs == sorted([100, 125, 50, 25, 5])