Problem statement

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

Solution

The idea is to keep two heaps: one for the top half of our data and another is for down half of our data. If we have even size $2n$, then we will keep two heaps with size $n$ If we have odd size $2n+1$, then we will keep size of the small heap $n+1$ ans the size of large heap $n$.

When we have new element num, we always put it to small heap, and then normalize our heaps: remove biggest element from the small heap and put it to the large heap. After this operation we can be sure that we have the property that the largest element in small heap is smaller than smaller elements in large heap.

However after this step if we had $n, n$ elements, we will have $n, n+1$ elements, so we need to put one element from large heap to small heap.

Complexity

Time complexity is just $O(1)$ to get median and $O(\log n)$ to add number. Space complexity is $O(n)$ after $n$ operations.

Code

class MedianFinder:
    def __init__(self):
        self.small, self.large = [], [] 

    def addNum(self, num):
        heappush(self.small, -num)           
        heappush(self.large, -heappop(self.small))
        
        if len(self.small) < len(self.large):
            heappush(self.small, -heappop(self.large))
            
    def findMedian(self):
        if len(self.large) != len(self.small):
            return -self.small[0]                  
        else:
            return (self.large[0] - self.small[0]) / 2

Remark

There is also solution with $O(\log n)$ time complexities for both operations, if we use Sorted List