1. 程式人生 > >(演算法總結)堆排序的應用:尋找中位數

(演算法總結)堆排序的應用:尋找中位數

設計一個數據結構,可動態地維護一組資料,且支援如下操作:

(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();
}