*** Find Median from Data Stream (Leetcode 295)

Problem Link: https://leetcode.com/problems/find-median-from-data-stream/

import heapq
class MedianFinder:

    def __init__(self):
        
        # We use two heaps.
        # One to store small values(maxheap)
        # Another to store large values(minheap)
        self.small=[]
        self.large=[]
        

    def addNum(self, num: int) -> None:
        
        # First push it into small arr(max heap)
        heapq.heappush(self.small,-1*num)
        
        # Now check if the max value in maxheap
        # is greater than the small value in minheap
        if(len(self.small)>0 and len(self.large)>0 and -1*self.small[0]>self.large[0]):
            val=-1*heapq.heappop(self.small)
            heapq.heappush(self.large,val)
        
        # If diff b/w sizes of small and large is > 2
        if(len(self.small)-len(self.large)>=2):
            val=-1*heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        
        if(len(self.large)-len(self.small)>=2):
            val=heapq.heappop(self.large)
            heapq.heappush(self.small,-1*val)
        
        
    def findMedian(self) -> float:
        
        if(len(self.small)==len(self.large)):
            return (-1*(self.small[0])+self.large[0])/2
        elif(len(self.small)>len(self.large)):
            return -1*self.small[0]
        else:
            return self.large[0]

Last updated