295. 数据流的中位数

思路

这题巧妙地利用了堆和中位数的性质。我们可以用一个大根堆和一个小根堆来分别模拟数组的前半部分和后半部分。这样我们就可以以O(1)的时间复杂度获得中位数，O(logN)的时间维护堆。

我们约定左边的大根堆的元素个数大于等于右边的小根堆元素个数，且两边的元素个数差最多为1，这样就可以解决奇数个元素时取元素的问题。为了维护这个性质，我们需要在插入元素时对堆进行处理，如代码所示。

Code

class MedianFinder {

    PriorityQueue<Integer> l = new PriorityQueue<>((a, b) -> b - a);
    PriorityQueue<Integer> r = new PriorityQueue<>();

    public MedianFinder() {
        
    }
    
    //维持左比右大一，或者左和右相等
    public void addNum(int num) {
        if (l.size() == r.size()) {
            if (r.isEmpty()) {
                l.offer(num);
            } else if (num <= r.peek()) {
                l.offer(num);
            } else {
                l.offer(r.poll());
                r.offer(num);
            }
        } else {
            if (num >= l.peek()) {
                r.offer(num);
            } else {
                r.offer(l.poll());
                l.offer(num);
            }
        }
    }
    
    public double findMedian() {
        if (l.size() == r.size()) {
            return (double)(l.peek() + r.peek()) / 2.0;
        }
        return l.peek();
    }
}