随时获取数据流的中位数

题目

有一个源源不断往外吐出整数的数据流，假设你有足够的空间来保存吐出的数。请设计一个方法，这个方法可以随时取出之前吐出所有数的中位数

要求

1. 如果已经保存了吐出的 N 个数，那么任意时刻将一个新数加入的过程，其时间复杂度不超过 O (logN)
2. 取得中位数的过程，时间复杂度为 O (1)

思路

建立一个大根堆，一个小根堆。每次来的一个数，和大根堆的堆顶比较，如果小于大根堆的堆顶，就加入大根堆；如果大于大根堆的堆顶，就加入小根堆

同时还要满足这两个堆中的元素个数之差不能超过 2（即 < 2）。例如大根堆中的元素现在有 3 个，小根堆中的元素有 1 个，此时就需要把大根堆的堆顶弹出，放入小根堆中；反之也一样。注意：每次往堆中加入数的同时，也要调整堆的结构

如果吐出的数据个数为偶数，则中位数是两个堆的堆顶相加除以 2；为奇数，中位数是元素个数较多的那个堆的堆顶

往堆里加入一个数的时间复杂度是 O (logN)，取出中位数的时间复杂度是 O (1)，满足题目要求

代码

import java.util.Comparator;
import java.util.PriorityQueue;

public class GetMedian {

    public static class MinHeapComparator implements Comparator<Integer> {

        @Override
        public int compare(Integer a, Integer b) {
            return a > b ? 1 : -1;
        }
    }

    public static class MaxHeapComparator implements Comparator<Integer> {

        @Override
        public int compare(Integer a, Integer b) {
            return b > a ? 1 : -1;
        }
    }

    public static void main(String[] args) {
        int[] arr = { 8, 3, 4, 6, 9, 7, 1 }; // 1 3 4 6 7 8 9
        double median = getMedian(arr);
        System.out.println(median);
    }

    private static double getMedian(int[] arr) {
        PriorityQueue<Integer> big = new PriorityQueue<Integer>(new MaxHeapComparator());
        PriorityQueue<Integer> small = new PriorityQueue<Integer>(new MinHeapComparator());
        for (int i = 0; i < arr.length; i++) {
            int cur = arr[i];
            if (big.isEmpty() || cur < big.peek())
                big.add(cur);
            else
                small.add(cur);
            if (big.size() - small.size() >= 2)
                small.add(big.poll());
            else if (small.size() - big.size() >= 2)
                big.add(small.poll());
        }
        if (arr.length % 2 == 0)
            return ((double) (big.peek() + small.peek()) / 2);
        else
            return big.size() > small.size() ? big.peek() : small.peek();
    }
}