题目描述
中位数是有序列表中间的数。如果列表长度是偶数,中位数则是中间两个数的平均值。
例如,
[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 范围内,你将如何优化你的算法?
算法
(堆) $O(\log n)$
中位数能把数据分成数量相等的两边,左边都小于等于中位数,右边都大于等于中位数。那么对应的,假如我们有两组数,左边组的数都小于等于右边组的数,那么左边组的最大值和右边组最小值的平均值就是当前数据的中位数。所以我们可以用两个堆来维护数据的左右半边,左边用大根堆维护,右边用小根堆维护,每次插入一个数的时候我们看它属于左半边还是右半边,如果当前大根堆为空或者它小于等于当前大根堆的堆顶那么它就属于左半边,否则它就属于右半边。另外为了维持左右两边数的数量,我们可以让大根堆的大小最多比小根堆大1
,小根堆的大小要严格小于或等于大根堆。这样当数据个数是奇数的时候中位数就是大根堆堆顶,是偶数的时候就是大根堆与小根堆堆顶的平均值。
对于两个follow-up
,如果数据全在[0, 100]
的范围内,我们可以开一个长度为101
的数组,利用计数排序 (counting sort) 的思想,记录每个数出现的次数还有数的总个数,返回中位数的时候计算一下需要的是第几个数然后可以遍历一遍数组来得到中位数。
如果 $99\%$ 的数都在[0, 100]
的范围内,我们对于小于0
和大于100
这两种情况分别开一个桶来记录出现的所有数并保持桶内有序,这里可以用平衡树或者有序数组来实现,插入这两个桶的时候我们需要维护桶内有序的性质,对于[0, 100]
之间的跟第一个follow-up
一样,同样的返回中位数的时候我们需要判断一下中位数落在哪个范围。
时间复杂度
插入一个数的时间复杂度为 $O(\log n)$,取中位数的时间复杂度为常数 $O(1)$。
C++ 代码
class MedianFinder {
public:
/** initialize your data structure here. */
priority_queue<int> maxHeap;
priority_queue<int, vector<int>, greater<int>> minHeap;
MedianFinder() {
}
void addNum(int num) {
if (maxHeap.empty() || num <= maxHeap.top())
maxHeap.push(num);
else
minHeap.push(num);
if (maxHeap.size() > minHeap.size() + 1)
{
minHeap.push(maxHeap.top());
maxHeap.pop();
}
else if (minHeap.size() > maxHeap.size())
{
maxHeap.push(minHeap.top());
minHeap.pop();
}
}
double findMedian() {
if (minHeap.size() == maxHeap.size())
return (minHeap.top() + maxHeap.top()) / 2.0;
return maxHeap.top();
}
};
/**
* Your MedianFinder object will be instantiated and called as such:
* MedianFinder* obj = new MedianFinder();
* obj->addNum(num);
* double param_2 = obj->findMedian();
*/