[LeetCode]（面试题41）数据流中的中位数

题目

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

例如，
[2,3,4] 的中位数是 3
[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

• void addNum(int num) - 从数据流中添加一个整数到数据结构中。
• double findMedian() - 返回目前所有元素的中位数。

示例 1：

输入：
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出：[null,null,null,1.50000,null,2.00000]

示例 2：

输入：
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出：[null,null,2.00000,null,2.50000]

限制：

• 最多会对 addNum、findMedia进行 50000 次调用。

解题思路

中位数可以把数据分为两部分，前一部分用大根堆，后一部分用小根堆，至于分成的两部分中的数字是否有序，我们并不关心。由此可以得到四个性质：

1. 大根堆中的所有数字都小于或等于大根堆的堆顶元素；
2. 小根堆中的所有数字都大于或等于小根堆的堆顶元素；
3. 大根堆的堆顶元素小于或等于小根堆的堆顶元素；
4. 大根堆中的所有数字都小于等于小根堆中的所有数字。

具体做法就是：

1. 当数据个数为偶数的时候，让两个堆中的元素相等，则两个堆顶元素的平均值就是中位数；
2. 当数据个数为奇数的时候，让大根堆的元素个数比小根堆的元素个数多1个，则大根堆的堆顶元素就是前一部分的最大值，也是整个数据流的中位数。

至于怎么平衡堆？设在还未插入新数之前，整个数据流的个数为 n，新插入的数为 x：

• 1）n 为奇数，那么在插入 x 之后，就变成了偶数。在插入之前，大根堆 A 的元素个数比小根堆 B 的元素个数多1，按道理需要将 x 插入小根堆 B，但 如果 x 小于大根堆 A 中的某些数，那么就将它先插入大根堆 A，再把大根堆 A 中最大的数插入小根堆 B，相当于做了一次置换操作。
• 2）n为偶数，那么在插入 x 之后，就变成了奇数。在插入之前，大根堆 A 的元素个数与小根堆 B 的元素个数相同，按道理需要将 x 插入大根堆 A ，但如果 x 大于小根堆 B 中的某些数，那么就将它先插入小根堆 B，再把小根堆 B 中最小的数插入大根堆 A，相当于做了一次置换操作。

算法流程：
设元素总数为 N = m + n，其中 m 和 n 分别为 大根堆 A 和小根堆 B 中的元素个数。
addNum(num) 函数：

• 1）当 m = n（即 N 为 偶数）：需向 A 添加一个元素。实现方法：将新元素 x 插入至 B ，再将 B 堆顶元素插入至 A ；
• 2）当 m != n（即 N 为 奇数）：需向 B 添加一个元素。实现方法：将新元素 x 插入至 A ，再将 A 堆顶元素插入至 B ；
  （注：假设插入数字 x 遇到情况1，由于 x 可能属于 “较大的一半” （即属于 ），因此不能将 x 直接插入至 A 。为了避免这种情况，统一先将 x 插入至 B ，再将 B 堆顶元素插入至 A 。这样就可以始终保持 A 保存较小一半、 B保存较大一半。情况2同理。）

findMedian() 函数：

• 1）当 m = n（N 为 偶数）：则中位数为 ( A 的堆顶元素 + B 的堆顶元素 ) / 2；
• 2）当 m != n（ N 为 奇数）：则中位数为 A 的堆顶元素。

复杂度分析：
时间复杂度：
查找中位数：O(1) ，获取堆顶元素使用 O(1) 时间。
添加数字：O(logN) ，堆的插入和弹出操作使用 O(logN) 时间。
空间复杂度：O(N)，其中 N 为数据流中的元素数量，大根堆 A 和小根堆 B 最多同时保存 N 个元素。

代码

class MedianFinder {
     PriorityQueue<Integer> A;
     PriorityQueue<Integer> B;
    /** initialize your data structure here. */
    public MedianFinder() {
        A = new PriorityQueue<>((a,b) -> b-a); // 大根堆
        B = new PriorityQueue<>((a,b) -> a-b); // 小根堆
    }
    
    public void addNum(int num) {
        int m = A.size();
        int n = B.size();
        if(m==n){
            B.add(num);
            A.add(B.poll());
        }else{
            A.add(num);
            B.add(A.poll());
        }
    }
    
    public double findMedian() {
        int m = A.size();
        int n = B.size();
        if(m==n){
            return (double)(A.peek()+B.peek())/2;
        }else{
            return (double)A.peek();
        }
    }
}

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