【剑指Offer-树】数据流中的中位数

题目描述

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

解题思路

由于是数据流，也就是说输入是动态变化的，所以数据结构最好选用动态数组，ArrayList比较合适。在计算中位数前对集合进行排序，然后返回中位数即可。

代码实现

public class Solution {

    private List<Integer> data = new ArrayList<>();

    public void Insert(int num) {
        data.add(num);
    }

    public Double GetMedian() {
        Collections.sort(data);
        System.out.println(data);
        if ((data.size() & 1) == 1) {
            return data.get(data.size() / 2).doubleValue();
        }else {
            Double median = (data.get(data.size() / 2 - 1).doubleValue() + data.get(data.size() / 2).doubleValue())/ 2.0;
            return median;
        }
    }
}

优化方法

上述方法虽然简单，但是对于每一次计算中位数都需要将数组进行排序，因而时间复杂度较高，即 n * nlog(n)。
那么，有什么办法可以使得排序的时间复杂度变低呢？我们想到了堆这个数据结构。当数据流中的数为偶数个时，入最小堆，但是先入最大堆，将最大堆中堆顶的数据取出，放入最小堆中。当数据流中的数据为奇数个时，入最大堆，但是先入最小堆，将最小堆堆顶的数据取出，放入最大堆中。
这样，最大堆中的数据都比最小堆中的数据小，并且最大堆中堆顶元素与最小堆堆顶元素是排序后数据中间相邻的两个数。当数据流中的数据为奇数个时，直接取最小堆堆顶的元素，为偶数个时，取最大堆与最小堆堆顶元素的平均值。

代码实现

public class OptimizeSolution {

    private PriorityQueue<Integer> minHeap = new PriorityQueue<>();// 小顶堆
    private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(15, new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });
    private int count = 0;

    public void Insert(int num) {
        // 偶数个数
        if ((count & 1) == 0) {
            // 先入大根堆
            maxHeap.offer(num);
            // 获取大根堆中最大值
            Integer maxNum = maxHeap.poll();
            // 将大根堆中的最大值放入小根堆
            minHeap.offer(maxNum);
        } else {
            // 奇数时，先入小根堆
            minHeap.offer(num);
            // 获取小根堆中的最小值
            Integer minNum = minHeap.poll();
            // 将这个最小值放入大根堆
            maxHeap.offer(minNum);
        }
        count++;
    }

    public Double GetMedian() {
        if ((count & 1) == 0) {
            int min = minHeap.peek();
            int max = maxHeap.peek();
            return (min + max) / 2.0;
        } else {
            return minHeap.peek() / 1.0;
        }
    }

    public static void main(String[] args) {
        OptimizeSolution s = new OptimizeSolution();
        s.Insert(5);
        s.Insert(3);
        s.Insert(2);
        s.Insert(1);
        System.out.println(s.GetMedian());
    }
}