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