leetcodealgorithms-templates4-heaps
 
 
 
import heapq
 
class Median:
    def __init__(self):
        self.small, self.large = [], []
 
    def insert(self, num):
        # Push to the max heap and swap with min heap if needed.
        heapq.heappush(self.small, -1 * num)
        if (self.small and self.large and (-1 * self.small[0]) > self.large[0]):
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        
        # Handle uneven size
        if len(self.small) > len(self.large) + 1:
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -1 * val)
        
    def getMedian(self):
        if len(self.small) > len(self.large):
            return -1 * self.small[0]
        elif len(self.large) > len(self.small):
            return self.large[0]
        
        # Even # of elements, return avg of two middle nums
        return (-1 * self.small[0] + self.large[0]) / 2