【leetcode 优先队列（堆） C++】295. Find Median from Data Stream

295. Find Median from Data Stream

class MedianFinder {
public:
    /** initialize your data structure here. */

    // 简单讲，把所有元素分成大的一半和小的一半，分别用小顶堆和大顶堆存储，两个堆顶元素就是中间的数

    int total_num = 0;  // 记录总的数量
    priority_queue<int, vector<int>, less<int> > PQ_less;  // 大顶堆，记录小的一半
    priority_queue<int, vector<int>, greater<int> > PQ_greater;  // 小顶堆，记录大的一半
    MedianFinder() {
    }
    
    void addNum(int num) {
        PQ_less.push(num);  // 向大顶堆中放数据
        if(total_num % 2) {  // 如果之前总的个数为奇数，就需要想小顶堆中移动一个数据
            PQ_greater.push(PQ_less.top());
            PQ_less.pop();
        }
        else if(total_num >= 2 && PQ_less.top() > PQ_greater.top()) {  // 如果大顶堆堆顶大于小顶堆堆顶，则需要交换数据
            int A = PQ_less.top(), B = PQ_greater.top();
            PQ_less.pop();
            PQ_greater.pop();
            PQ_less.push(B);
            PQ_greater.push(A);
        }
        total_num++;
    }
    
    double findMedian() {
        if(total_num % 2 == 1) return PQ_less.top();  // 奇数个时，直接返回大顶堆堆顶
        return 0.5 * (PQ_less.top() + PQ_greater.top());  // 偶数时，计算两个堆的平均
    }
};

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