【JZ-41】数据流中的中位数（堆）

题目

参考

算法思路

重点是每次插入新的数据后都要使得列表中的数保持有序，然后返回中间元素。为了使得时间复杂度尽可能小，可以分别建立小根堆和大根堆，前者保存列表中较大的一半数据，后者保存列表中较小的一半数据，这样要获得中位数直接根据两个堆的堆顶数据进行计算即可。

算法流程

假设共有 N 个数据，小根堆 A 中存放 m 个，大根堆 B 中存放 n 个。当 N 为偶数时，m = n = N/2；当 N 为奇数时，m = (N + 1)/2，n = (N - 1)/2

addNum(int num)：

1. 若 N 为偶数，应向 A 添加一个元素：
   将 num 插入 B，再将 B 堆顶元素插入 A，从而保持 A 始终存放较大的一半， B 始终存放较小的一半。
2. 若 N 为奇数，应向 B 添加一个元素：
   将 num 插入 A，再将 A 堆顶元素插入 B

findMedian()：

设 A 的堆顶元素为 $a$，B 的堆顶元素为 $b$，

1. 若 N 为偶数，中位数为 $\frac{a+b}{2}$
2. 若 N 为奇数，中位数为 $a$

具体代码

class MedianFinder {

    /** initialize your data structure here. */
    Queue<Integer> A, B;
    public MedianFinder() {
        A = new PriorityQueue<>();//小根堆，保存较大的一半
        B = new PriorityQueue<>(new Comparator<Integer>(){
            @Override
            public int compare(Integer o1, Integer o2){
                return o2.compareTo(o1);
            }
        });//大根堆，保存较小的一半
    }
    
    public void addNum(int num) {
        //若N为偶数，向A添加一个元素
        if(A.size() == B.size()){
            B.offer(num);
            A.offer(B.poll());
        }else{//否则向B添加一个元素
            A.offer(num);
            B.offer(A.poll());
        }
    }
    
    public double findMedian() {
        return A.size() == B.size() ? (A.peek() + B.peek()) / 2.0 : A.peek();
    }
}

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

关于大根堆自定义比较器：
java的lamda表达式，方法参数是（x,y），方法体是 return y-x 。因为方法体只有一行return，这种情况下可以简化，省略return
即，可以简写为B = new PriorityQueue<>((x, y) -> (y - x));

复杂度分析

• 时间复杂度：$O(logN)$，添加数据，即堆的插入和删除的时间开销为 $log(N)$，查找中位数的时间开销为 $O(1)$
• 空间复杂度：$O(N)$