Learning On the Way

295. Find Median from Data Stream

Design a data structure that supports the following two operations:

• void addNum(int num) - Add a integer number from the data stream to the data structure.
• double findMedian() - Return the median of all elements so far.

from heapq import heappop, heappush 

class MedianFinder: 
    def __init__(self): 
        """ 
        initialize your data structure here. 
        """ 
        self._low =  [] # maxheap, allow to have one extra element
        self._high = [] # minheap 
        self._median = None
    
    def low_to_high(self):
        heappush(self._high, -heappop(self._low))

    def high_to_low(self):
        heappush(self._low, -heappop(self._high))

    def balance(self):
        if (len(self._low) > len(self._high) + 1):
            self.low_to_high()
        elif len(self._high) > len(self._low):
            self.high_to_low()

    def addNum(self, num: int) -> None: 
        heappush(self._low, -num)
        if len(self._high) > 0 and -self._low[0] > self._high[0]: 
            self.low_to_high()
        self.balance()

    def findMedian(self) -> float: 
        if len(self._low) > len(self._high): 
            return -self._low[0] 
        else: 
            return (-self._low[0] + self._high[0]) / 2