剑指 Offer 41. 数据流中的中位数

发布时间 2023-09-13 22:14:29作者: luxiayuai
class MedianFinder {
public:
    /** initialize your data structure here. */
    // 注意小根堆的定义方式
    priority_queue<int, vector<int>, greater<int>> up; // 小根堆,默认放从大到小后一半的数
    priority_queue<int> down; // 大根堆,默认放从小到大排序前一半的数
    // 注意:大根堆的堆顶小于小根堆的堆顶,两个数是大小是紧挨着的,两个队呈现两个三角对顶的关系
    MedianFinder() {

    }
    
    void addNum(int num) {
        // 如果下面是空的,或者num小于大根堆堆顶元素
        // 这里一定要注意加上小根堆为空的条件,否则当判断num <= down.top()时,由于down中不存在数据,会越界访问
        if(down.empty() || num <= down.top()) {
            // 将num查到大根堆
            down.push(num);

            // 当插入一个新数后,大根堆比小根堆多两个数,此时大根堆需要将堆顶pop掉作为小根堆的堆顶
            if(down.size() > up.size() + 1) {
                up.push(down.top());
                down.pop();
            } 
        } 
        else {
            // 如果num大于大根堆堆顶,并且小根堆不为空
            up.push(num);
            // 如果大根堆比小根堆的数多了,那么将小根堆的堆顶pop掉作为大根堆的堆顶
            if(down.size() < up.size()) {
                down.push(up.top());
                up.pop();
            }
        }
    }
    
    double findMedian() {
        // 这时要比较两个堆存数的个数。正常情况下两种情况:
        // 第一种,大根堆比小根堆多一个数,那么中位数就是大根堆的堆顶
        // 第二种,大根堆和小根堆存数的数量相等,那么中位数就是大根堆和小根堆的均值
        // 因此,每新插入一个数,都要维护大小根堆的size的动态平衡。

        // 如果两个堆的数据个数相加为奇数,说明大根堆比小根堆多一个数,中位数是大根堆堆顶
        if((down.size() + up.size()) % 2) return down.top();

        // 如果数据个数相等,中位数就是两者堆顶的平均数
        // 注意,除以2.0,否则结果会下取整,变为整形
        else return (down.top() + up.top()) / 2.0;
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
 */