LeetCode 295. Find Median from Data Stream（优先队列）

295. Find Median from Data Stream

Hard

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

For example,
[2,3,4], the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

void addNum(int num) - Add a integer number from the data stream to the data structure.
double findMedian() - Return the median of all elements so far.

Example:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3)
findMedian() -> 2

题意

编写一个类，支持插入数据addNum和寻找中位findMedian数操作

思路

这个问题比较开放，根据addNum和findMedian操作次数的不同，应设计不同的数据结构。试了一下测试用例里addNum和findMedian操作的数量是1:1，因此采用双优先队列法。
具体来说，就是维护一个最大优先队列lower存储比较小的那一半元素，维护一个最小优先队列higher存储比较大的那一半元素，在addNum过程中，时刻保持0 <= lower.size() - higher.size() <=1 && lower.peek() <= higher.peek()。当数据总量是奇数时，中位数就是lower.peek()，当数据总量是偶数时，中位数是(lower.peek() + higher.peek())/2.
插入数据addNum的时间复杂度就是优先队列插入的时间复杂度O(logn)，找中位数的时间复杂度O(logn)，空间复杂度O(n).

代码

class MedianFinder {
    private int length = 0;
    private PriorityQueue<Integer> lower = new PriorityQueue<Integer>(new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });
    private PriorityQueue<Integer> higher = new PriorityQueue<Integer>();

    /** initialize your data structure here. */
    public MedianFinder() {
        
    }
    
    /**
    * O(log(n))
    */
    public void addNum(int num) {
        if (length == 0) {
            lower.add(num);
            length = 1;
            return;
        }
        if (lower.size() <= higher.size()) {
            if (num <= higher.peek()) {
                lower.add(num);
            } else {
                lower.add(higher.poll());
                higher.add(num);
            }
        } else {
            if (num >= lower.peek()) {
                higher.add(num);
            } else {
                higher.add(lower.poll());
                lower.add(num);
            }
        }
        ++length;
        // System.out.println(lower);
        // System.out.println(higher);
    }
    
    /**
    * O(1)
    */
    public double findMedian() {
        if ((length & 1) == 0) {
            return (double) (lower.peek() + higher.peek()) / 2;
        } else {
            return lower.peek();
        }
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */