C++的 STL堆 實作擷取中位數

前言

堆資料結構 使用的是優先級隊列實作，建立堆的時候需要指定堆中元素的排列方式，即最大堆或者最小堆

最大堆即 堆頂元素為堆中最大的元素

最小堆即 堆頂元素為堆中最小堆元素

如下為一個最大堆

中位數：

一組數排序後，如果元素個數如下

奇數個數n：(int) n/2 的數

偶數個數n: (int) n/2 和(int) n/2 +1的平均數

同樣此時想要擷取一組數列中的中位數，且同樣使用時間複雜度為O(n)進行解決，這個時候使用堆的資料結構是最為有效的

解決辦法如下：

動态維護一個最大堆與一個最小堆，最大堆存儲一半資料，最小堆存儲 一般資料，維持最大堆的堆頂比最小堆的堆頂小。

在維護這兩個堆的過程中核心為

1. 最大堆的堆頂元素需要小于堆小堆的堆頂元素
2. 兩個堆元素個數相差不能超過1

保證核心的情況下需要考慮如下三種情況：

a. 最大堆的元素和最小堆的元素個數相等

此時入堆需要根據進入元素和兩個堆頂元素大小比較的結果進而選擇入哪個堆

b. 最大堆的元素個數大于最小堆的元素個數

此時入堆需要分兩種情況：進入元素的大小小于最大堆的堆頂，進入元素大小大于最大堆的堆頂

c. 最大堆的元素個數小于最小堆的元素個數

此時入堆需要分兩種情況：進入元素的大小大于最小堆的堆頂，進入元素大小小于最小堆的堆頂

實作過程如下(文末有完整測試代碼)：

/*
核心：
1. 保證最大堆的堆頂小于最小堆的堆頂，這樣就保證了最大堆的元素都小于最小堆的元素
2. 兩個堆中堆元素個數相差不能超過1
*/
void GetMedian::push(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 { //否push最小堆
            small_heap.push(num);
        }
    } else if(big_heap.size() > small_heap.size()) {//第二種情況
        /*
        此時入堆元素小于最大堆的堆頂，按照第一個核心，需要入最大堆
        但是為了保證兩個堆大小個數相差不能超過1
        則需要将最大堆堆堆頂取出來放入最小堆，再将目前元素入堆
        */
        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);
        }
    }
}      

#include <iostream>
#include <queue>

using namespace std;

class GetMedian{
    private:
    priority_queue<int,vector<int>,greater<int>> small_heap;
    priority_queue<int,vector<int>,less<int>> big_heap;
    public:
    GetMedian(){};
    void push(int num);
    double getMedia(); 
};

void GetMedian::push(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 GetMedian::getMedia() {
    double median;

    if (small_heap.size() == big_heap.size()) {
        median = ((double)small_heap.top() + (double)big_heap.top()) / 2; 
    } else if (small_heap.size() < big_heap.size()) {
        median = (double)big_heap.top();
    } else {
        median = (double)small_heap.top();
    }

    return median;
}

int main() {
    GetMedian m;
    int tmp;
    int n;

    cout << "input the number of element " << endl;
    cin >> n;

    cout << "add " << n << " element \n" << endl;
    for (int i =0;i < n; ++i) {
        cin >> tmp;
        m.push(tmp);
    }

    cout << "median is " << m.getMedia() << endl;
    return 0;
}      

input the number of element 
6
add 6 element 

3 4 2 1 5 6
median is 3.5


input the number of element 
5
add 5 element 

6 5 7 3 1
median is 5