def test_pop(self): h = MaxHeap(elements=to_add, do_copy=True) self.assertEqual(h.peek(), max(to_add)) self.assertEqual(len(h), len(to_add)) for i, v in enumerate(reversed(sorted(to_add))): print(i, v) self.assertEqual(h.pop(), v) self.assertEqual(len(h), len(to_add) - i - 1) self.assertEqual(len(h), 0) self.assertIsNone(h.pop()) for i in range(len(to_add_2)): h.insert(to_add_2[i]) self.assertEqual(h.peek(), max(to_add_2[:i + 1])) self.assertEqual(len(h), i + 1)
class PriorityQueueHeap: "Implementation of a priority queue as a heap" # Reuse of elements used in tree implementation is inefficient # because we don't need a linked list. Oh well. def __init__(self, vals=None, priorities=None): if priorities is not None and vals is not None: if len(priorities) != len(vals): raise ValueError("Must be as many priorities as values.") elements = [ PriorityQueueElement(p, v) for v, p in zip(vals, priorities) ] else: elements = None self._h = MaxHeap(elements=elements) def insert(self, value, priority): element = PriorityQueueElement(priority, value) self._h.insert(element) def pop_max(self): element = self._h.pop() if element is None: return element return element.pop() def peek(self): element = self._h.peek() if element is None: return element return element.peek() def is_empty(self): return len(self) < 1 def shrink_to_size(self): "Shrinks heap capacity to current size" self._h.shrink_to_size() def __len__(self): return len(self._h)
def test_insert(self): h = MaxHeap() for i in range(len(to_add)): h.insert(to_add[i]) self.assertEqual(h.peek(), max(to_add[:i + 1])) self.assertEqual(len(h), i + 1)