class Median(object): def __init__(self, lst=[]): self.minHeap = PriorityQueue() self.maxHeap = MaxHeap() self.size = len(lst) if lst: lst = sorted(lst) mid = len(lst) // 2 self.minHeap = PriorityQueue(lst[mid:]) self.maxHeap = MaxHeap(lst[0:mid]) def get_median(self): if self.size == 0: raise IndexError("empty heap") if self.size == 1 or self.size % 2 != 0: return self.minHeap.get_min()[0] else: return self.maxHeap.get_max()[0] def insert(self, prio, ele): if self.size == 0: self.minHeap.insert(prio, ele) elif self.size == 1: self.minHeap.insert(prio, ele) prio, ele = self.minHeap.delete_min() self.maxHeap.insert(prio, ele) else: min_in_minHeap = self.minHeap.get_min()[0] max_in_maxHeap = self.maxHeap.get_max()[1] if prio >= min_in_minHeap: self.minHeap.insert(prio, ele) else: self.maxHeap.insert(prio, ele) if not self.is_balanced(): self.balance() self.size += 1 def is_balanced(self): return self.minHeap.size() <= self.maxHeap.size() + 1 and \ self.minHeap.size() >= self.maxHeap.size() def balance(self): while not self.is_balanced(): if self.minHeap.size() >= self.maxHeap.size() + 1: prio, ele = self.minHeap.delete_min() self.maxHeap.insert(prio, ele) else: prio, ele = self.maxHeap.delete_max() self.minHeap.insert(prio, ele)
class MaxHeap(object): def __init__(self, lst=[]): lst = [(-prio, ele) for (prio, ele) in lst] self.maxHeap = PriorityQueue(lst) def insert(self, priority, ele): self.maxHeap.insert(-priority, ele) def delete_max(self): prio, ele = self.maxHeap.delete_min() return -prio, ele def get_max(self): prio, ele = self.maxHeap.get_min() return -prio, ele def size(self): return self.maxHeap.size()