剑指 Offer 41题：数据流中的中位数

1.题目描述

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

例如，

[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]

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof

2.思路分析

建立一个 小顶堆 A 和 大顶堆 B ，各保存列表的一半元素，且规定：

• A保存较大的一半元素，长度为N / 2（N为偶数）或 (N + 1) / 2（N为奇数）
• B保存较小的一半元素，长度为N / 2（N为偶数）或 (N + 1) / 2（N为奇数）

随后，中位数可仅根据 A,B 的堆顶元素计算得到。 

时间复杂度：

• 查找中位数 O(1) ： 获取堆顶元素使用 O(1) 时间；
• 添加数字 O(log N) ： 堆的插入和弹出操作使用 O(logN) 时间。

空间复杂度 O(N)： 其中 N 为数据流中的元素数量，小顶堆 A 和大顶堆 B 最多同时保存 N 个元素。

3.解答 

class MedianFinder {
    //堆

    //存储元素
    PriorityQueue<Integer> small, big;

    /** initialize your data structure here. */
    public MedianFinder() {
        small = new PriorityQueue<>(); //小顶堆，存储较大的一部分
        big = new PriorityQueue<>((x, y) -> (y - x)); //大顶堆，存储较小的一部分
    }
    
    //添加元素
    public void addNum(int num) {
        if(small.size() == big.size()){
            //如果两个堆存储的数量相等，则将新元素存储到大顶堆中
            small.offer(num);
            big.offer(small.poll());
        }else{
            //否则将新元素存到samll中
            big.offer(num);
            small.offer(big.poll());
        }
    }
    
    //返回中位数
    public double findMedian() {
        //判断数组元素的数量
        if(small.size() != big.size()){
            //如果是奇数，中位数就是big堆的堆顶元素
            return big.peek();
        }else{
            //是偶数
            return (small.peek() + big.peek()) / 2.0;
        }
    }
}