Find Median from Data Stream — Two-Heap Approach

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Design a data structure that supports:

  • addNum(num) — add a number from the stream
  • findMedian() — return the current median

For even count: median = average of two middle elements.


Key Insight — Two Heaps

  • lo = max-heap (lower half) — Python: negate for max-heap
  • hi = min-heap (upper half)
  • Invariant: len(lo) == len(hi) or len(lo) == len(hi) + 1
lo (max-heap): [5, 3, 1]  hi (min-heap): [7, 9]
Median = lo.top() = 5   (odd total)

After each add: rebalance so |len(lo)-len(hi)| <= 1 and lo.max <= hi.min.


Solutions

Python

import heapq

class MedianFinder:
    def __init__(self):
        self.lo = []
        self.hi = []

    def addNum(self, num: int) -> None:
        heapq.heappush(self.lo, -num)
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        if len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

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

JavaScript

class MedianFinder {
    constructor() { this.lo=[]; this.hi=[]; }
    addNum(n) {
        this._pushMax(this.lo, n);
        this._pushMin(this.hi, -this._popMax(this.lo));
        if (this.hi.length > this.lo.length) this._pushMax(this.lo, -this._popMin(this.hi));
    }
    findMedian() {
        return this.lo.length > this.hi.length ? -this.lo[0] : (-this.lo[0] + this.hi[0]) / 2;
    }
    _pushMax(h,v){h.push(-v); let i=h.length-1; while(i>0){let p=(i-1)>>1; if(h[p]<h[i]){[h[p],h[i]]=[h[i],h[p]];i=p;}else break;}}
    _popMax(h){[h[0],h[h.length-1]]=[h[h.length-1],h[0]]; const v=h.pop(); let i=0; while(true){let l=2*i+1,r=2*i+2,m=i; if(l<h.length&&h[l]>h[m])m=l; if(r<h.length&&h[r]>h[m])m=r; if(m===i)break;[h[i],h[m]]=[h[m],h[i]];i=m;} return -v;}
    _pushMin(h,v){h.push(v); let i=h.length-1; while(i>0){let p=(i-1)>>1; if(h[p]>h[i]){[h[p],h[i]]=[h[i],h[p]];i=p;}else break;}}
    _popMin(h){[h[0],h[h.length-1]]=[h[h.length-1],h[0]]; const v=h.pop(); let i=0; while(true){let l=2*i+1,r=2*i+2,m=i; if(l<h.length&&h[l]<h[m])m=l; if(r<h.length&&h[r]<h[m])m=r; if(m===i)break;[h[i],h[m]]=[h[m],h[i]];i=m;} return v;}
}

Java

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

C++

#include <queue>
using namespace std;
class MedianFinder {
    priority_queue<int> lo;
    priority_queue<int,vector<int>,greater<int>> hi;
public:
    void addNum(int n) {
        lo.push(n); hi.push(lo.top()); lo.pop();
        if (hi.size()>lo.size()) { lo.push(hi.top()); hi.pop(); }
    }
    double findMedian() {
        return lo.size()>hi.size() ? lo.top() : (lo.top()+hi.top())/2.0;
    }
};

C

/* C: two binary heap arrays, manual push/pop with heapify */

Complexity

OperationTimeSpace
addNumO(log n)O(n)
findMedianO(1)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro