Two Heaps Priority Queue For Medians

MedianPQ.py

class MedianPQ:
    def __init__(self):
        self.left = []  # left part of an array (max heap)
        self.right = []  # right part of an array (min heap)
        self.removed = collections.defaultdict(int)  # lazy removal
        self.balance = 0

    def add(self, val):
        if not self.left or val <= -self.left[0]:
            heapq.heappush(self.left, -val)
            self.balance -= 1
        else:
            heapq.heappush(self.right, val)
            self.balance += 1
        self.rebalance()

    def remove(self, val):
        self.removed[val] += 1
        if val <= -self.left[0]:
            self.balance += 1
        else:
            self.balance -= 1
        self.rebalance()
        self.lazy_remove()

    def rebalance(self):
        while self.balance < 0:
            heapq.heappush(self.right, -heapq.heappop(self.left))
            self.balance += 2
        while self.balance > 0:
            heapq.heappush(self.left, -heapq.heappop(self.right))
            self.balance -= 2

    def lazy_remove(self):
        while self.left and (-self.left[0] in self.removed):
            self.removed[-self.left[0]] -= 1
            if self.removed[-self.left[0]] == 0:
                del self.removed[-self.left[0]]
            heapq.heappop(self.left)
        while self.right and (self.right[0] in self.removed):
            self.removed[self.right[0]] -= 1
            if self.removed[self.right[0]] == 0:
                del self.removed[self.right[0]]
            heappop(self.right)

    def get_median(self):
        if self.balance == 0:
            return (-self.left[0] + self.right[0]) / 2
        elif self.balance < 0:
            return -self.left[0]
        else:
            return self.right[0]