295. 数据流的中位数
中位数是有序列表中间的数。如果列表长度是偶数,中位数则是中间两个数的平均值。
例如,
[2,3,4] 的中位数是 3
[2,3] 的中位数是 (2 + 3) / 2 = 2.5
设计一个支持以下两种操作的数据结构:
void addNum(int num)
- 从数据流中添加一个整数到数据结构中。double findMedian()
- 返回目前所有元素的中位数。
示例:
addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3)
findMedian() -> 2
进阶:
- 如果数据流中所有整数都在 0 到 100 范围内,你将如何优化你的算法?
- 如果数据流中 99% 的整数都在 0 到 100 范围内,你将如何优化你的算法?
本体思考的方式是对顶堆。尽管我们输入的数据流不是有序的,但是我们可以利用平衡树让他们有序。本题目其实并不需要平衡树,因为我们只维护有序数组中的中位数,因而只需要1-2个数据。这里我们使用2个堆来进行模拟。
有序数组的左侧我们看做一个大根堆,右侧看成一个小根堆,因为有序的特性,左侧的大根堆的顶部尽管是左侧最大的数字,但是其实也不会大于右侧的小根堆的顶部。明白了这个性质,我们接下来为了求出中位数,可以设置成以下2种情况
- 小根堆的长度等于大根堆的长度
- 大根堆的长度比小根堆的长度多一
也有别的设置方法。我们取大根堆-小根堆的差值为0或1的做法。
在我们添加数字的时候,有以下的情况
- 初始情况:也就是2个堆的size都等于0。因为奇数个数据的时候大根堆的长度是大于小根堆的,我们直接push到大根堆里。
- 2个堆的size相等而且非0。这时候我们根据提前设置好的情况,想让大根堆的数量+1的话,要考虑2个子情况。
该数字比小根堆的顶部小,这个情况很简单,那我们就继续维护有序的特性,正好左侧的大根堆还需要一个数字,直接插入即可。
该数字比小根堆的顶部大,对于这个情况,为了维护有序,同时保持大根堆的数量+1的话,我们则需要把小根堆现在最小的数字,放进大根堆里面去,再把新的数字放进小根堆里。 - 大根堆的size比小根堆的size多一个。这时候,我们需要对小根堆的数量+1,也要考虑2个子情况。
该数字比大根堆的顶部大,这个情况很简单,我们直接将该数字放进小根堆使得2个堆的size相等
该数字比大根堆的顶部小,这个情况,为了继续维护大根堆的有序,我们把左侧大根堆最大的数字,加进小根堆里来,同时将该数字放进小根堆里
做完了以上的准备工作,求median就可以在$O(1)$时间完成。
-
如果是奇数,此时大根堆的数量必定比小根堆多1,直接返回大根堆顶部。
-
如果是偶数,取2个顶部除2即可。
class MedianFinder {
public:
priority_queue<int> maxHeap;
priority_queue<int, vector<int>,greater<int>> minHeap;
MedianFinder() {
}
void addNum(int num) {
int m = maxHeap.size(), n = minHeap.size();
if (m == n) {
if (minHeap.empty() || num <= minHeap.top()) maxHeap.push(num); // 情况1和情况2.1
else { // 情况2.2
maxHeap.push(minHeap.top());
minHeap.pop();
minHeap.push(num);
}
} else {
if (num >= maxHeap.top()) {
minHeap.push(num); // 情况3.1
} else { // 情况3.2
minHeap.push(maxHeap.top());
maxHeap.pop();
maxHeap.push(num);
}
}
assert(abs((int)maxHeap.size() -(int)minHeap.size()) <= 1); // 检测是否满足size要求
}
double findMedian() {
if (maxHeap.size() == minHeap.size()) {
return double(maxHeap.top() + minHeap.top()) /2.0;
} else return maxHeap.top();
}
};