Find Median from Data Stream

Problem: Find Median from Data Stream

We can use two heaps, one to store small numbers and one to store big numbers. If the sizes of the two heaps differentiated by one, we need to maintain the sizes. Since the python use min-heaps, we need to store negative numbers in the heap so that the biggest number of small numbers are at top. Besides, we need to deal with situations with odd and even numbers.

Code in Python:

import heapq

class MedianFinder:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.lheap = []
        self.rheap = []


    def addNum(self, num):
        """
        Adds a num into the data structure.
        :type num: int
        :rtype: void
        """
        if not self.lheap:
            heapq.heappush(self.lheap, -num)
        elif num < -self.lheap[0]:
            heapq.heappush(self.lheap, -num)
        else:
            heapq.heappush(self.rheap, num)

        if len(self.lheap)-len(self.rheap)>1:
            heapq.heappush(self.rheap, -heapq.heappop(self.lheap))
        elif len(self.rheap)>len(self.lheap):
            heapq.heappush(self.lheap, -heapq.heappop(self.rheap))


    def findMedian(self):
        """
        Returns the median of current data stream
        :rtype: float
        """
        if not self.lheap: return None
        if (len(self.lheap)+len(self.rheap)) % 2: return -self.lheap[0]
        else: return (-self.lheap[0]+self.rheap[0])/2.0