今天做了一题堆相关的题目295数据流的中位数,题解巧妙地用一个大根堆和一个小根堆来快速计算中位数。
其中添加新数字的写法十分精妙:如果一个数字要添加到小根堆,就先添加到大根堆,再将大根堆堆顶的元素转移至小根堆;反之亦然,以此能够保证小根堆中的元素永远大于大根堆中的元素。
其中,C++优先级队列底层默认是一个大根堆(使用了less<xxx>,更换成greater<xxx>,即变为小根堆)。
class MedianFinder {
public:
/** initialize your data structure here. */
MedianFinder() {
}
void addNum(int num) {
if(que1.size()>=que2.size()){
que1.push(num);
que2.push(que1.top());
que1.pop();
}else{
que2.push(num);
que1.push(que2.top());
que2.pop();
}
}
double findMedian() {
if(que1.size()>que2.size())
return que1.top();
if(que1.size()==que2.size())
return (double)(que1.top()+que2.top())/2;
return que2.top();
}
private:
priority_queue<int,vector<int>,less<int>>que1; //大根堆
priority_queue<int,vector<int>,greater<int>>que2; //小根堆
};
/**
* Your MedianFinder object will be instantiated and called as such:
* MedianFinder* obj = new MedianFinder();
* obj->addNum(num);
* double param_2 = obj->findMedian();
*/