【LeetCode】数据流的中位数 [H]（堆）

295. 数据流的中位数 - 力扣（LeetCode）

一、题目

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。

例如 arr = [2,3,4] 的中位数是 3 。
例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。
实现 MedianFinder 类:

1. MedianFinder() 初始化 MedianFinder 对象。
2. void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
3. double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受。

示例 1：

输入
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
输出
[null, null, null, 1.5, null, 2.0]

解释
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

提示：

• -10^5 <= num <= 10^5
• 在调用 findMedian 之前，数据结构中至少有一个元素
• 最多 5 * 10^4 次调用 addNum 和 findMedian

二、代码

class MedianFinder {
    // 大根堆和小根堆
    public PriorityQueue<Integer> maxHeap;
    public PriorityQueue<Integer> minHeap;

    public MedianFinder() {
        this.maxHeap = new PriorityQueue<>((a, b) -> b - a);
        this.minHeap = new PriorityQueue<>((a, b) -> a - b);
    }
    
    public void addNum(int num) {
        // 这里要记得先判断一下大根堆是不是空，如果是空的话说明此时两个堆一定都是空，将这个数直接加入到大根堆中
        if (maxHeap.isEmpty() || maxHeap.peek() >= num) {
            maxHeap.add(num);
        } else {
            minHeap.add(num);
        }

        // 调平两个堆容量
        if (maxHeap.size() - minHeap.size() >= 2) {
            minHeap.add(maxHeap.poll());
        } else if (minHeap.size() - maxHeap.size() >= 2) {
            maxHeap.add(minHeap.poll());
        }
    }
    
    public double findMedian() {
        // 两个堆一样多，堆顶相加除以2就是中位数
        if (maxHeap.size() == minHeap.size()) {
            return  (maxHeap.peek() + minHeap.peek()) / 2D;
        // 两个堆不一样多，说明此时数据个数为奇数，那种size较大的那个堆的堆顶就是中位数    
        } else {
            return maxHeap.size() > minHeap.size() ? maxHeap.peek() : minHeap.peek();
        }
    }
}

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

三、解题思路 

用两个堆，大根堆 + 小根堆。

在大根堆里是由大到小排列的，在小根堆里是由小到大排列的

对于大根堆以及对于小根堆来说加入一个数字或者弹出堆顶都是logN级别

在加入数据的时候满足如下原则：

1. 当前数<=大根堆项，入大根堆
2. 当前数>大根堆项，入小根堆
3. 当大根堆size - 小根堆size大于等于2，就将大根堆堆顶转移到小根堆
4. 当小根堆size - 大根堆size大于等于2，就将小根堆堆顶转移到大根堆

所有数里较小的一半在大根堆里，所有数较大的一半在小根堆里。并且因为是小根堆和大根堆结构，所以两个结构的堆顶一定是较小的那一半数中的最大值，较大的那一半数中的最小值，因此大根堆的顶部和小根堆的顶部一定能够算出为中位数。 

计算中位数的时候：

1. 如果此时两个堆中的数据总数为奇数（换句话说就是两个堆的size不一样大，因为结构保证了两个堆的大小差不会大于等于2），则存储数据较多的那个堆的堆顶就是中位数。
2. 如果此时两个堆中的数据总数为偶数（换句话说就是两个堆的size一样大），中位数就是两个堆的堆顶相加除以2。