CodeEasy v0.1

Find Median from Data Stream

Hard Subject: Heap / Priority Queue
Time Complexity
O(log N) add, O(1) find
Space Complexity
O(N)

Problem Description

The median is the middle value in an ordered integer list. Implement the MedianFinder class that can add integers to a data stream and return the median at any time.

Optimal Solution

Python 
import heapq

class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (store negatives)
        self.large = []  # min-heap

    def addNum(self, num):
        heapq.heappush(self.small, -num)
        # balance: every element in small <= every element in large
        if self.small and self.large and (-self.small[0]) > self.large[0]:
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        # size balance
        if len(self.small) > len(self.large) + 1:
            val = -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, -val)

    def findMedian(self):
        if len(self.small) > len(self.large):
            return -self.small[0]
        elif len(self.large) > len(self.small):
            return self.large[0]
        return (-self.small[0] + self.large[0]) / 2.0