【Leetcode 热题 100】295. 数据流的中位数

问题背景

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

• 例如 $arr=[2,3,4]$ 的中位数是 $3$。
• 例如 $arr=[2,3]$ 的中位数是 $(2+3)/2=2.5$。
  实现 MedianFinder 类：
• MedianFinder() 初始化 MedianFinder 对象。
• void addNum(int num) 将数据流中的整数 $num$ 添加到数据结构中。
• double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 $10^{-5}$ 以内的答案将被接受。

数据约束

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

解题过程

这题分类讨论的思路比较麻烦，但是实现起来相当简洁，积累一下为好。
整体的结构是由一个最大堆和一个最小堆构成的，维护结构的过程中让这两个堆中的元素始终相等或最大堆中的元素比最小堆中的元素多，并且相差不能超过 $1$。
最大堆中存储数据流中较小的那部分元素，最小堆中存储数据流中较大的那部分元素，这时中位数就可以根据两个堆顶元素方便地计算出来。数字数量不等时，中位数就是最大堆的堆顶元素；元素数量相等时，中位数就是两个堆顶元素的平均值。
以下 $num$ 表示数据流中的新元素，$max$ 表示最小堆中的最大值，$min$ 表示最大堆中的最小值。

• 如果两个堆中元素数量相等：
  • 当前元素较大的情况下， $num$ 应被添加到最小堆中。这样一来就不符合定义了（最小堆比最大堆元素数量多$1$），将 $max$ 调整到最大堆中。
  • 当前元素较小的情况下， $num$ 应被添加到最大堆中。这种情况同样可以先将元素添加到最小堆中，再将 $max$ 调整到最大堆中。
• 如果两个堆中元素数量不等：
  • 当前元素较小的情况下， $num$ 应被添加到最大堆中。这样一来就不符合定义了（最大堆比最小堆元素数量多$2$），将 $min$ 调整到最小堆中。
  • 当前元素较大的情况下， $num$ 应被添加到最小堆中。这种情况同样可以先将元素添加到最大堆中，再将 $min$ 调整到最小堆中。

实际写的时候，只要搞清楚什么情况下该怎么倒腾元素就不会出错。
Java 中的类默认自带一个空参构造器，所以实际上 MedianFinder() 可以不写。

具体实现

class MedianFinder {
    private final PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);
    private final PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    
    public void addNum(int num) {
        if(maxHeap.size() == minHeap.size()) {
            minHeap.offer(num);
            maxHeap.offer(minHeap.poll());
        } else {
            maxHeap.offer(num);
            minHeap.offer(maxHeap.poll());
        }
    }
    
    public double findMedian() {
        if(maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        }
        return (minHeap.peek() + maxHeap.peek()) / 2.0;
    }
}

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