Leetcode 295. Find Median from Data Stream

Two heap to save the smaller half and larger half number respectively. 

We need to keep heap as balanced while insertion, i.e., the different of two heaps' size must be <= 1.

If we find the insertion will make the different larger than 1, we move the head number of a heap to the other, and then insert the number into the heap so that heap are kept balanced. 

findMeadian is obvious. Only need to consider the top numbers of two heap.

public class MedianFinder {
    // a min heap to save the larger half numbers
    private PriorityQueue<Integer> minHeap;
    // a max heap to save the smaller half numbers
    private PriorityQueue<Integer> maxHeap;
    
    /** initialize your data structure here. */
    public MedianFinder() {
        minHeap = new PriorityQueue<Integer>();
        maxHeap = new PriorityQueue<Integer>(
            new Comparator<Integer>() {
                public int compare(Integer a, Integer b) {
                    return b - a;
                }
            });
    }
    
    public void addNum(int num) {
        // the size of two heaps are same
        if (minHeap.size() == 0 && maxHeap.size() == 0) {
            minHeap.offer(num);
        } else if (minHeap.size() > maxHeap.size()){
            // heap's size is different, keep balance of two heaps
            if (num > minHeap.peek()) {
                // move the top of the minHeap to the maxHeap
                maxHeap.offer(minHeap.poll());
                // add num to the minHeap
                minHeap.offer(num);
            } else {
                maxHeap.offer(num);
            }
        } else if (minHeap.size() < maxHeap.size()){
            if (num < maxHeap.peek()) {
                minHeap.offer(maxHeap.poll());
                maxHeap.offer(num);
            } else {
                minHeap.offer(num);
            }
        } else {
            if (num > minHeap.peek()) {
                minHeap.offer(num);
            } else {
                maxHeap.offer(num);
            }
        }
    }
    
    public double findMedian() {
        if (minHeap.size() == maxHeap.size()) {
            return (minHeap.peek()+maxHeap.peek())/2.0;
        } else if (minHeap.size() > maxHeap.size()) {
            return (double) minHeap.peek();
        }
        return (double) maxHeap.peek();
    }
}