Find Median from Data Stream — Two Heaps

Problem 298 · Find Median from Data Stream

Difficulty: Hard · Pattern: Two Heaps (lower max + upper min)

Intuition

Maintain:

lo: max-heap of smaller half
hi: min-heap of larger half

After each add, balance so len(lo) == len(hi) or len(lo) == len(hi)+1.

Median: if equal sizes, average top of both; else top of lo.

Solutions

# Python
import heapq
class MedianFinder:
    def __init__(self):
        self.lo = []  # max-heap (negate)
        self.hi = []  # min-heap

    def addNum(self, num):
        heapq.heappush(self.lo, -num)
        # Balance: move max of lo to hi
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        # Ensure lo has >= elements
        if len(self.lo) < len(self.hi):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

    def findMedian(self):
        if len(self.lo) > len(self.hi):
            return -self.lo[0]
        return (-self.lo[0] + self.hi[0]) / 2

// Java
class MedianFinder {
    PriorityQueue<Integer> lo = new PriorityQueue<>(Collections.reverseOrder());
    PriorityQueue<Integer> hi = new PriorityQueue<>();
    public void addNum(int num) {
        lo.offer(num);
        hi.offer(lo.poll());
        if (lo.size() < hi.size()) lo.offer(hi.poll());
    }
    public double findMedian() {
        return lo.size() > hi.size() ? lo.peek() : (lo.peek() + hi.peek()) / 2.0;
    }
}

// C++
class MedianFinder {
    priority_queue<int> lo;
    priority_queue<int,vector<int>,greater<int>> hi;
public:
    void addNum(int num) {
        lo.push(num); hi.push(lo.top()); lo.pop();
        if (lo.size()<hi.size()) { lo.push(hi.top()); hi.pop(); }
    }
    double findMedian() {
        return lo.size()>hi.size() ? lo.top() : (lo.top()+hi.top())/2.0;
    }
};

Complexity

addNum: O(log n)
findMedian: O(1)