def test_max_heap_heapify():
heap = MaxHeap([5, 5, 4, 4, 3])
assert heap.max() == 5
assert heap.pop_max() == 5
assert heap.pop_max() == 5
assert heap.max() == 4
assert heap.pop_max() == 4
heap = MaxHeap([4, 3, 5, 4, 5])
assert heap.max() == 5
assert heap.pop_max() == 5
assert heap.pop_max() == 5
assert heap.max() == 4
assert heap.pop_max() == 4
def test_max_heap_pop():
heap_max = MaxHeap()
try:
heap_max.pop_max()
except:
pass
else:
raise AssertionError
heap_max.push(1)
heap_max.push(3)
assert heap_max.pop_max() == 3
assert heap_max.max() == 1
assert heap_max.pop_max() == 1
try:
heap_max.pop_max()
except:
pass
else:
raise AssertionError
heap_max.push(9)
heap_max.push(8)
heap_max.push(5)
assert heap_max.pop_max() == 9
assert heap_max.pop_max() == 8
assert heap_max.pop_max() == 5
def test_max_heap_max():
heap_max = MaxHeap()
try:
heap_max.max()
except:
pass
else:
raise AssertionError
heap_max.push(1)
assert heap_max.max() == 1
heap_max.push(2)
assert heap_max.max() == 2
heap_max.push(0)
assert heap_max.max() == 2
class PrefixWithMedian:
def __init__(self):
self.max_heap = MaxHeap()
self.min_heap = MinHeap()
def insert(self, val) -> None:
len_prefix = len(self.max_heap) + len(self.min_heap)
if len_prefix <= 1:
self.min_heap.push(val)
else:
ref_val = self.min_heap.min()
if len_prefix % 2 == 0:
self._insert_and_shift_median(val, ref_val)
else:
self._insert_and_keep_median(val, ref_val)
def get_median(self):
return self.min_heap.min()
def get_prefix(self):
raise NotImplemented
def _insert_and_shift_median(self, new_val, cur_median) -> None:
"""
Вставка элемента если индекс медианного элемента изменится
Те
len(prefix) % 2 = 0
"""
if new_val < cur_median:
self.max_heap.push(new_val)
else:
self.min_heap.pop_min()
self.max_heap.push(cur_median)
self.min_heap.push(new_val)
def _insert_and_keep_median(self, new_val, cur_median) -> None:
"""
Вставка элемента если индекс медианного элемента не изменится
Те
len(prefix) % 2 = 1
"""
if new_val > cur_median:
self.min_heap.push(new_val)
else:
if self.max_heap.max() < new_val:
self.min_heap.push(new_val)
else:
last_max_heap = self.max_heap.pop_max()
self.max_heap.push(new_val)
self.min_heap.push(last_max_heap)