295 find median from data stream

https://leetcode.com/problems/find-median-from-data-stream/

class MedianFinder {
    PriorityQueue<Integer> loPq;
    PriorityQueue<Integer> hiPq;
    public MedianFinder() {
        loPq = new PriorityQueue<>(new Comparator<Integer>(){
            @Override
            public int compare(Integer a, Integer b) {
                return -1 * a.compareTo(b);
            }
        });
        hiPq = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        if (hiPq.isEmpty()) {
            loPq.offer(num);
        } else {
            int lo = loPq.peek();
            int hi = hiPq.peek();
            if (num < hi) {
                loPq.offer(num);
            } else {
                hiPq.offer(num);
            }
        }
        int length = loPq.size() + hiPq.size();
        if (length % 2 == 0) {
            if (loPq.size() > hiPq.size()) {
                int mid = loPq.poll();
                hiPq.offer(mid);
            } else if (loPq.size() < hiPq.size()) {
                int mid = hiPq.poll();
                loPq.offer(mid);
            }
            
        } else if (length % 2 == 1) {
            if (loPq.size() > hiPq.size() + 1) {
                int mid = loPq.poll();
                hiPq.offer(mid);
            } else if (loPq.size() < hiPq.size() + 1) {
                int mid = hiPq.poll();
                loPq.offer(mid);
            }
        }
    }
    
    public double findMedian() {
        int length = loPq.size() + hiPq.size();
        if (hiPq.size() == 0) {
            return loPq.peek();
        }
        if (length % 2 == 0) {
            return ((double)loPq.peek() + hiPq.peek()) / 2;
        } else {
            return Math.min(loPq.peek(), hiPq.peek());
        }
    }
}

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

用一个max heap记录上半部分,用一个min heap记录下半部分,注意如果数字过大要放到先办部分的minheap里面,同时注意两半部分size的平衡

Last updated