【优先队列 堆】剑指 Offer 41. 数据流中的中位数

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

将数据分为两半，较小的一半用大顶堆存放，较大的一半用小顶堆存放。当前数据总数为 $cnt$ ，这样当一个新的数据 $num$ 进来时，考虑下面两种情况：

• $cnt\%2==0$ ：先将 $num$ 压入小顶堆，然后将小顶堆堆顶元素压入大顶堆，最后 $pop$ 掉小顶堆堆顶元素；

  这种情况时，两个堆内的元素数量相同，那么加入新元素后，前半段数据会比后半段数据多1个，先将新元素和后半段元素比较，然后将后半段元素中最小的放到前半段来。

• $cnt\%2==1$ ：先将 $num$ 压入大顶堆，然后将大顶堆堆顶元素压入小顶堆，最后 $pop$ 掉大顶堆堆顶元素；

  这种情况时，前半段数据数量比后半段多1，那么加入新元素后，两段数据数量应该相同，也就是需要往后半段添加元素。先将新元素和前半段元素比较，然后将前半段元素中最大的放到后半段来。

经过上面的压入操作，保证了两个堆顶元素就是处在中间位置的元素，如果当前总数为奇数个，中位数即为大顶堆的堆顶元素，如果当前总数为偶数个，中位数即为两个堆顶元素的平均数。

class MedianFinder {
    priority_queue<int, vector<int>, less<int>> bigger; // 大顶堆
    priority_queue<int, vector<int>, greater<int>> lesser; // 小顶堆
    int cnt = 0;	// 数据总量
public:
    /** initialize your data structure here. */
    MedianFinder() {}
    
    void addNum(int num) {
        if (cnt % 2 == 0) {
            lesser.push(num);
            bigger.push(lesser.top());
            lesser.pop();
        } else {
            bigger.push(num);
            lesser.push(bigger.top());
            bigger.pop();
        }
        cnt++;
    }
    
    double findMedian() {
        double ret = 0.0;
        if (cnt % 2 == 1) {
            ret = bigger.top();
        } else {
            ret = (bigger.top() + lesser.top()) / 2.0;
        }
        return ret;
    }
};