数据流中的中位数
题目链接:https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof/
思路:维护一个最大堆和一个最小堆,最大堆总是存储比中位数小的一半数据,最小堆总是存储比中位数大的一半数据,从而实现O(1)返回中位数。需要注意维护一个invariant:两个堆的size要么相等要么差1,因此在插入新数据时:
- 若此时两个堆size相等,先插入到最小堆,然后将最小堆的栈顶数据pop出来插入到最大堆
- 若此时两个堆size差1(必然是最大堆比最小堆多1个),则先插入到最大堆,然后将最大堆的栈顶数据pop出来插入到最小堆
上述过程能够保证两个堆总是存放正确的数据。
参考代码:
class MedianFinder {
public:
/** initialize your data structure here. */
MedianFinder() {
}
void addNum(int num) {
if (maxHeap.size() == minHeap.size()) {
minHeap.push(num);
maxHeap.push(minHeap.top());
minHeap.pop();
return;
}
maxHeap.push(num);
minHeap.push(maxHeap.top());
maxHeap.pop();
}
double findMedian() {
if (maxHeap.size() == minHeap.size()) {
return (maxHeap.top() + minHeap.top())/2.0;
}
return maxHeap.top();
}
private:
priority_queue<int, vector<int>, less<int>> maxHeap;
priority_queue<int, vector<int>, greater<int>> minHeap;
};