题目描述

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。我们使用Insert()方法读取数据流,使用GetMedian()方法获取当前读取数据的中位数。

解题思路:

假设我们已经有数字 1 2 3 4 5 6 7 8,那么我们可以把前半部分的4个数字做一个大堆big,后4个数字做一个小堆small。然后(big.top()+small.top())/2即为答案。这两个堆我们允许它们之间的容量相差为1。同时small的最小值大于big的最大值。

很自然地,我们想到的维护方法是:每次insert一个数字时候,我们检查big.top(),假如比big.top()大的话,big.insert(num),假如比small.top()小的话,small.insert(num),否则我们可以在两个堆中的容量较小的堆insert. 注意每次insert后,我们都需要检查两个堆的大小,假如有一个堆比另一个堆的size大2,我们就从大size的堆把数字pop出来给小size的堆。

class Solution {
public:
priority_queue<int,vector<int>,greater<int>> pq_small;
priority_queue<int,vector<int> , less<int>> pq_big;
void Insert(int num)
{
if(!pq_small.size()){pq_small.push(num);return;}
if(!pq_big.size()){
if(pq_small.top()<num){
pq_big.push(pq_small.top());
pq_small.pop();
pq_small.push(num);
}else pq_big.push(num);
return;
}
if(num>pq_small.top()){
pq_small.push(num);
if(pq_small.size()>pq_big.size()+1){
pq_big.push(pq_small.top());
pq_big.pop();
}
return;
}
if(num<pq_big.top()){
pq_big.push(num);
if(pq_big.size()>pq_small.size()+1){
pq_big.push(pq_small.top());
pq_small.pop();
}
return;
}
if(pq_small.size()<pq_big.size())pq_small.push(num);
else pq_big.push(num);
}

double GetMedian()
{
int tot = pq_small.size()+pq_big.size();
if(tot%2==0)return ((double)pq_small.top()+(double)pq_big.top()) / 2;
else{
if(pq_small.size()>pq_big.size())return pq_small.top();
else return pq_big.top();
}
}
};