(演算法總結)堆排序的應用:尋找中位數
設計一個數據結構,可動態地維護一組資料,且支援如下操作:
(1)新增元素:void addNum(int num)
(2)返回這組資料中的中位數 double findMedian()
【思考】如何獲取一組元素的中位數
(1)首先,我們馬上想到的方法,最直觀的方法就是:新增元素的同時進行排序操作(直插sort)addNum的複雜度是O(n),findMedian的複雜度則是O(1);
(2)我們也可以考慮在查詢中位數的時候進行排序,這樣addNum的複雜度是O(1),而findMedian的複雜度則是O(nlogn)。然而,如果addNum和findMedian的操作都是隨機操作,共進行n次,則整體的時間複雜度最佳為O(n2)
這樣的複雜度很明顯不能滿足要求,因此,我們需要試圖尋找最佳的解決策略,這裡我們巧妙地利用了堆的性質:
這裡,我們動態地維護了一個大頂堆 max_heap 和一個小頂堆 min_heap,這兩個堆各自存放“一半”的資料(這裡的“一半”是指兩個堆的size()相差小於等於1),且維持 max_heap 的堆頂 <= min_heap 的堆頂。如何實現呢?我們這樣設計:
在插入新元素x時,如果max_heap為空,直接將x壓入max_heap;
否則,先比較 max_heap.size() 與 min_heap.size():
若max_heap與min_heap元素個數相同:
若 x < max_heap 的堆頂,將x壓入 max_heap;
否則,將x壓入 min_heap;
若 max_heap 比 min_heap 的size小(需要通過調整兩個堆,以保證平衡):
若 x < max_heap 的堆頂,直接將x壓入 max_heap;
否則,將 min_heap 的堆頂彈出並壓入 max_heap,再將x壓入min_heap
若 max_heap 比 min_heap 的size大(同樣需要通過調整兩個堆):
若 x > max_heap 的堆頂,直接將x壓入 min_heap;
否則,將 max_heap 的堆頂彈出並壓入 min_heap,再將x壓入max_heap
至此addNum()的流程已經完成,接下來需要考慮如何求中位數了。
若max_heap與min_heap元素個數相同:
median = max_heap 的堆頂與 min_heap 的堆頂的算數平均值
若 max_heap 比 min_heap 的size小:
median = min_heap 的堆頂
否則:
median = max_heap 的堆頂
程式碼如下:
class MedianFinder {
public:
void addNum(int num);
double find_median();
private:
std::priority queue<int, std::vector<int>, std::less<int>> max_heap;
std::priority queue<int, std::vector<int>, std::greater<int>> min_heap;
};
void MedianFinder::addNum(int num) {
if (max_heap.empty()) {
max_heap.push(num);
return;
}
if (max_heap.size() == min_heap.size()) {
if (num < max_heap.top()) {
max_heap.push(num);
} else {
min_heap.push(num);
}
} else if (max_heap.size() < min_heap.size()) {
if (num < max_heap.top()) {
max_heap.push(num);
} else {
max_heap.push(min_heap.top());
min_heap.pop();
min_heap.push(num);
}
} else {
if (num < max_heap.top()) {
min_heap.push(max_heap.top());
max_heap.pop();
max_heap.push(num);
} else {
min_heap.push(num);
}
}
}
int MedianFinder::find_median() {
if (max_heap.size() == min_heap.size()) {
return (max_heap.top() + min_heap.top()) / 2.0;
} else if (max_heap.size() < min_heap.size()) {
return min_heap.top();
}
return max_heap.top();
}