剑指 Offer 41. 数据流中的中位数

剑指 Offer 41. 数据流中的中位数

class MedianFinder {
    PriorityQueue<Integer> pq1, pq2;

    /** initialize your data structure here. */
    public MedianFinder() {
        pq1 = new PriorityQueue<>((a, b) -> Integer.compare(a, b));  // 返回最小的元素，存放较大的一半
        pq2 = new PriorityQueue<>((a, b) -> Integer.compare(b, a));  // 返回最大的元素，存放较小的一半
    }
    
    public void addNum(int num) {
        if(pq1.size() == pq2.size()){
            pq1.offer(num);
            pq2.offer(pq1.poll());
        }else{
            pq2.offer(num);
            pq1.offer(pq2.poll());
        }
    }
    
    public double findMedian() {
        if(pq1.size() == pq2.size()){
            return (pq1.peek() + pq2.peek()) / 2.0;
        }else{
            return pq2.peek();
        }
    }
}