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)