每次进行addNum后都sort一次会TLE
使用一个最大堆和最小堆
【STL】使用priority_queue构造堆
每次 addNum操作都改变堆一次
最大堆的size大于最小堆时，最大堆的堆顶元素就是中位数
最大堆的size小于最小堆时，最小堆的堆顶元素就是中位数
二者相等时，最大最小堆顶元素的平均值是中位数
C++代码：
#include <algorithm>
#include <iostream>
#include <queue>
using namespace std;
class MedianFinder
{
public:
priority_queue<int, vector<int>, greater<int>> small_heap;
priority_queue<int, vector<int>, less<int>> big_heap;
/** initialize your data structure here. */
MedianFinder()
{
}
void addNum(int num)
{
if (big_heap.empty())
{
big_heap.push(num);
return;
}
if (big_heap.size() == small_heap.size())
{
if (num < big_heap.top())
big_heap.push(num);
else
small_heap.push(num);
}
else if (big_heap.size() > small_heap.size())
{
if (num <= big_heap.top())
{
small_heap.push(big_heap.top());
big_heap.pop();
big_heap.push(num);
}
else
small_heap.push(num);
}
else
{
if (num >= small_heap.top())
{
big_heap.push(small_heap.top());
small_heap.pop();
small_heap.push(num);
}
else
big_heap.push(num);
}
}
double findMedian()
{
if (big_heap.size() == small_heap.size())
return (double(big_heap.top()) + double(small_heap.top()) / 2);
else if (big_heap.size() > small_heap.size())
return double(big_heap.top());
else
return double(small_heap.top());
}
};
int main()
{
MedianFinder s;
s.addNum(1);
cout << s.findMedian() << endl;
s.addNum(2);
cout << s.findMedian() << endl;
s.addNum(3);
cout << s.findMedian() << endl;
return 0;
}
