class TestMaxHeap(unittest.TestCase):
def setUp(self):
self.heap = MaxHeap()
def test_basic_initialization_and_repr(self):
self.assertEqual(repr(self.heap), '[]')
def test_insert(self):
self.heap.insert(4)
self.assertEqual(repr(self.heap), '[4]')
self.assertEqual(self.heap.size, 1)
self.heap.insert(4)
self.assertEqual(repr(self.heap), '[4, 4]')
self.assertEqual(self.heap.size, 2)
self.heap.insert(6)
self.assertEqual(repr(self.heap), '[6, 4, 4]')
self.assertEqual(self.heap.size, 3)
self.heap.insert(1)
self.assertEqual(repr(self.heap), '[6, 4, 4, 1]')
self.assertEqual(self.heap.size, 4)
self.heap.insert(7)
self.assertEqual(repr(self.heap), '[7, 6, 4, 1, 4]')
self.assertEqual(self.heap.size, 5)
def test_get_max(self):
self.assertEqual(self.heap.get_max(), None)
self.heap.insert(-7)
self.assertEqual(self.heap.get_max(), -7)
self.heap.insert(7)
self.assertEqual(self.heap.get_max(), 7)
self.heap.insert(5)
self.assertEqual(self.heap.get_max(), 7)
self.heap.insert(12)
self.assertEqual(self.heap.get_max(), 12)
def test_extract_min(self):
self.heap.insert(4)
self.heap.insert(5)
self.heap.insert(7)
self.heap.insert(2)
self.heap.insert(-1)
self.assertEqual(self.heap.extract_max(), 7)
self.assertEqual(self.heap.extract_max(), 5)
self.assertEqual(self.heap.extract_max(), 4)
self.assertEqual(self.heap.extract_max(), 2)
self.assertEqual(self.heap.extract_max(), -1)
self.assertEqual(self.heap.extract_max(), None)
def test_build_heap(self):
self.heap.build_heap([4, 4, 6, 1, 7])
self.assertEqual(repr(self.heap), '[7, 4, 6, 1, 4]')
class Median:
def __init__(self):
self.h_low = MaxHeap()
self.h_high = MinHeap()
def add_element(self, value):
if self.h_low.heap_size == 0 or value < self.h_low.max():
self.h_low.insert(value)
if self.h_low.heap_size - self.h_high.heap_size > 1:
self.h_high.insert(self.h_low.extract_max())
else:
self.h_high.insert(value)
if self.h_high.heap_size - self.h_low.heap_size > 1:
self.h_low.insert(self.h_high.extract_min())
def get_median(self):
if (self.h_low.heap_size + self.h_high.heap_size) % 2 == 0:
return self.h_low.max(), self.h_high.min()
else:
if self.h_low.heap_size > self.h_high.heap_size:
return self.h_low.max()
else:
return self.h_high.min()
def get_maxheap_elements(self):
return self.h_low.heap
def get_minheap_elements(self):
return self.h_high.heap
def heap_sort(lst):
"""
This functions is sorting a list according to heap sort in O(nlog(n))
:param lst: The list to sort
:return: The sorted list
"""
size = len(lst)
heap = MaxHeap(lst)
sorted_lst = []
for i in range(size):
sorted_lst.append(heap.extract_max())
return sorted_lst[::-1]
class MedianMaintenance:
def __init__(self):
self.hlow_heap = MaxHeap()
self.hhigh_heap = MinHeap()
def compute_median(self, i):
self.insert_heap(i)
self.balance_heap()
return self.median()
def balance_heap(self):
if self.hhigh_heap.size - self.hlow_heap.size > 1 : # rebalance heap to keep it balanced
high = self.hhigh_heap.extract_min()
self.hlow_heap.insert(high)
elif self.hlow_heap.size - self.hhigh_heap.size > 1:
low = self.hlow_heap.extract_max()
self.hhigh_heap.insert(low)
def insert_heap(self, i):
if self.hlow_heap.is_empty():
low = None
else:
low = self.hlow_heap.peek_max()
if self.hhigh_heap.is_empty():
high = None
else:
high = self.hhigh_heap.peek_min()
if low is None or i < low:
self.hlow_heap.insert(i)
elif high is not None and i > high:
self.hhigh_heap.insert(i)
else:# i wedged inbetween insert in first heap by default
self.hlow_heap.insert(i)
def median(self):
if self.hhigh_heap.size - self.hlow_heap.size == 1:
return self.hhigh_heap.peek_min()
else:# default choice when hlow is bigger/same size as hhigh
return self.hlow_heap.peek_max()
def test_extract_max(self):
array = [2, 3, 3, 7, 10, 11, 42]
max_heap = MaxHeap(values=array)
self.assertEqual(max_heap.extract_max(), 42)
