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

题目链接

维护两个优先队列

class MedianFinder {

    Queue<Integer> A, B;
    /** initialize your data structure here. */
    public MedianFinder() {
        // A保存较大的一半，小顶堆
        A = new PriorityQueue<>();
        // B保存较小的一半，大顶堆
        B = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer a, Integer b) {
                return b.compareTo(a);
            }
        });
    }
    
    public void addNum(int num) {
        // 当A.size() = B.size()时，要向A里面加元素，加了以后总数为奇数，中位数为A顶
        // 当A.size() > B.size()时，要向B里面加元素，加了以后总数为偶数，中位数为AB顶的平均
        // 向A中加元素时，要先加到B，B再将顶给A，如果直接加A，可能违反AB规则（A所有元素 >= B）
        // 向B中加元素同理
        if (A.size() == B.size()) {
            B.offer(num);
            A.offer(B.poll());
        }
        else  {
            A.offer(num);
            B.offer(A.poll());
        }
    }
    
    public double findMedian() {
        if (A.size() == B.size()) {
            return (A.peek() + B.peek()) / 2.0;
        }
        else {
            return A.peek();
        }
    }
}

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