剑指63.数据流中的中位数

题目描述

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。我们使用Insert()方法读取数据流，使用GetMedian()方法获取当前读取数据的中位数。

 

思路

如果能够保证数据容器左边的数据都小于右边的数据，那么即使左、右两边内部的数据没有排序，也可以根据左边最大的数及右边最小的数得到中位数。

因此，用一个最大堆实现左边的数据容器，用一个最小堆实现右边的数据容器。往堆中插入一个数据的时间效率是O(logn)。由于只需要O(1)的时间就可以得到位于堆顶的数据，因此得到中位数的时间复杂度是O(1)

首先测试一下PriorityQueue的默认实现是最大堆还是最小堆。

import java.util.PriorityQueue;
public class JZ63 {
    public static void main(String[] args) {
        //PriorityQueue默认实现是最小堆
        /*PriorityQueue<Integer> queue = new PriorityQueue<>();
        queue.add(5);
        queue.add(1);
        queue.add(7);
        System.out.println(queue.poll()); // 输出1
        */
        // 使用lambda表达式，改为最大堆
        PriorityQueue<Integer> queue = new PriorityQueue<>((o1, o2) -> o2-o1);
        queue.add(5);
        queue.add(7);
        queue.add(1);
        System.out.println(queue.poll()); // 输出7
    }
}

解法

import java.util.PriorityQueue;

public class Solution {
    private int count = 0; // 数据流中数的个数
    private PriorityQueue<Integer> queue1 = new PriorityQueue<>((o1, o2) -> (o2-o1)); // 最大堆，用来找左区间的最大值
    private PriorityQueue<Integer> queue2 = new PriorityQueue<>(); // 最小堆，用来找右区间的最小值
    //每次插入小顶堆的是当前大顶堆中最大的数
    //每次插入大顶堆的是当前小顶堆中最小的数
    //这样保证小顶堆中的数永远大于等于大顶堆中的数
    //中位数就可以方便地从两者的根结点中获取了
    public void Insert(Integer num) {
        if ((count & 1) == 0){ // 数据流中数的个数是偶数，放到左区间，相当于count % 2 == 0
            /*if (!queue2.isEmpty() && num > queue2.peek()){
                queue2.offer(num);
                num = queue2.poll();
            }
            queue1.offer(num);*/
            queue2.offer(num);
            int min = queue2.poll();
            queue1.offer(min);
        }else{ // 奇数时，下一个数字放入右区间
            /*if (!queue1.isEmpty() && num < queue1.peek()){
                queue1.offer(num);
                num = queue1.poll();
            }
            queue2.offer(num);*/
            queue1.offer(num);
            int max = queue1.poll();
            queue2.offer(max);
        }
        count++; // 要记得维护数据流中数的个数！！
    }
    public Double GetMedian() {
        if ((count & 1) == 0){ // 偶数
            return (queue1.peek() + queue2.peek()) / 2.0;
        }else{
            return (double) queue1.peek();
        }
    }
}