Find Median from Data Stream — Two-Heap Approach
Advertisement
Problem
Design a data structure that supports:
addNum(num)— add a number from the streamfindMedian()— 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-heaphi= min-heap (upper half)- Invariant:
len(lo) == len(hi)orlen(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
| Operation | Time | Space |
|---|---|---|
| addNum | O(log n) | O(n) |
| findMedian | O(1) | — |
Advertisement