算法通关村第十四关:黄金挑战-数据流的中位数

黄金挑战-数据流的中位数

1.数据流中中位数的问题

LeetCode295
https://leetcode.cn/problems/find-median-from-data-stream/

思路分析

中位数的问题，我们一般都可以用 大顶堆+小顶堆 来求解

小顶堆（minHeap）：存储所有元素中较大的一半，堆顶元素是其中最小的数
大顶堆（maxHeap）：存储所有元素中较小的一半，堆顶元素是其中最大的数

相当于，把所有元素分成了大和小两半，计算中位数，只需要大的那半的最小值和小的那半的最大值即可。

如[1,2,3,4,5]，分成小顶堆[3,4,5]，大顶堆[2,1]，中位数为3

过程
加1 小顶堆[1] 大顶堆[] 中位数 1
加2 小顶堆[2] 大顶堆[1] 中位数 (2+1)/2
加3 小顶堆[2,3] 大顶堆[1] 中位数 3
加4 小顶堆[3,4] 大顶堆[2,1] 中位数 (3+2)/2
加5 小顶堆[3,4,5] 大顶堆[2,1] 中位数 3

代码实现

Java中的堆（即优先队列）可方便实现，c++、python里没有提供堆结构，需要自己构造

class MedianFinder {
    // 小顶堆中存储的是比较大的元素，堆顶是其中的最小值
    PriorityQueue<Integer> minHeap;
    // 大顶堆中存储的是比较小的元素，堆顶是中期的最大值
    PriorityQueue<Integer> maxHeap;

    public MedianFinder() {
        this.minHeap = new PriorityQueue<>();
        this.maxHeap = new PriorityQueue<>((a, b) -> b - a);
    }
    
    public void addNum(int num) {
        // 小顶堆存储的是比较大的元素，num大于小顶堆中根元素时进入minHeap
        if (minHeap.isEmpty() || num > minHeap.peek()){
            minHeap.offer(num);
            // 如果minHeap比maxHeap多2个元素，就平衡一下
            if (minHeap.size() - maxHeap.size() > 1){
                maxHeap.offer(minHeap.poll());
            }
        }else{
            maxHeap.offer(num);
            // 这样可以保证多的那个元素肯定在minHeap
            if(maxHeap.size() - minHeap.size() > 0){
                minHeap.offer(maxHeap.poll());
            }
        }

    }
    
    public double findMedian() {
        if(minHeap.size() > maxHeap.size()){
            return minHeap.peek();
        }else if(minHeap.size() < maxHeap.size()){
            return maxHeap.peek();
        }else{
            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();
 */