算法 - 剑指Offer 数据流中的中位数

题目

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。
例如，
[2,3,4] 的中位数是 3
[2,3] 的中位数是 (2 + 3) / 2 = 2.5
设计一个支持以下两种操作的数据结构：
void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。

解题思路

这题比较难， 要用到新的知识： 大小顶堆，堆可以看成一个完全二叉树， 大顶堆定义是根节点的大小比其左右子节点都要大或者相等，小顶堆是根节点的大小比其左右子节点都要小或者相等， 大顶堆适合升序， 小顶堆适合降序。然后这里我们一整个数据流可以升序也可以降序， 我们这里选择降序，那么我们就可以用小顶堆 + 大顶堆来做， Java中可以用PriorityQueue来实现大小顶堆，只需要将数字插入，他们就会自动排序了，当数据流为奇数，大小顶堆大小不一样的时候， 我们将始终保持小顶堆比大顶堆多一个，这样取中位数的时候直接弹出小顶堆的最后一位即可， 当数据流为偶数的时候， 大小顶堆的大小是一样的， 这时候获取大小顶堆的最后一位去做一个平均值即可， 最后是这里的添加数字， 按照这里的想法， 当数据流为奇数，然后要插入一位数的时候， 我们要确保小顶堆的最小数值数据插入到大顶堆中， 这里直接add（PriorityQueue里面会自动排序）然后弹出最后一位，当数据流为偶数的时候， 我们相反， 需要从大顶堆弹出一位到小顶堆中， 具体代码入下。

Java代码实现

import java.util.*;
public class MedianFinder {
    Queue<Integer> bigQueue;
    Queue<Integer> smallQueue;
    /** initialize your data structure here. */
    public MedianFinder() {
        bigQueue = new PriorityQueue<>((o1,o2)->o2.compareTo(o1));
        smallQueue = new PriorityQueue<>((o1,o2)->o1.compareTo(o2));
    }
    public void addNum(int num) {
        if(smallQueue.size() == bigQueue.size()){
            bigQueue.add(num);
            smallQueue.add(bigQueue.poll());
        }else{
            smallQueue.add(num);
            bigQueue.add(smallQueue.poll());
        }
    }
    public double findMedian() {
        int isOdd = bigQueue.size() + smallQueue.size();
        if(isOdd%2 == 1){
            return smallQueue.peek();
        }else{
            return (bigQueue.peek() + smallQueue.peek()) / 2.0;
        }
    }
}