[LeetCode 295] Find Median from Data Stream

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.

Examples: 

[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.

For example:

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

solution:

 use a max heap on left side to represent elements that are less than effective median, and a min heap on right side to represent elements that are greater than effective median.

After processing an incoming element, the number of elements in heaps differ utmost by 1 element. When both heaps contain same number of elements, we pick average of heaps root data as effective median. When the heaps are not balanced, we select effective median from the root of heap containing more elements.

complexity is O(lgn)

class MedianFinder {
    PriorityQueue<Integer> minHeap = new PriorityQueue<Integer> (111, new Comparator<Integer>(){
        @Override
        public int compare(Integer a, Integer b) {
            return a-b;
        }
    });
    PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer> (111, new Comparator<Integer>(){
        @Override
        public int compare(Integer a, Integer b) {
            return b-a;
        }
    });
    // Adds a number into the data structure.
    public void addNum(int num) {
        if(minHeap.size()==0 && maxHeap.size()==0) {
            minHeap.add(num);
        }else if(maxHeap.size()==0 && num>minHeap.peek()){
            maxHeap.add(minHeap.poll());
            minHeap.add(num);
        }else if(num<minHeap.peek()){
            maxHeap.add(num);
        }else {
            minHeap.offer(num);
        }
        balance();
    }
    // balance both heap, the difference is <= 1
    public void balance() {
        int h1 = maxHeap.size();
        int h2 = minHeap.size();
        if(Math.abs(h1-h2)<=1) {
            return;
        }else {
            if(h1-h2>1) {
                minHeap.add(maxHeap.poll());
            }else {
                maxHeap.add(minHeap.poll());
            }
        }
    }
    // Returns the median of current data stream
    public double findMedian() {
        double res = 0.0;
        int h1 = maxHeap.size();
        int h2 = minHeap.size();        
        if(h1==0 && h2 ==0) return res;
        else if(h1==0 || h2 == 0){
            return h1==0? (double)minHeap.peek(): (double)maxHeap.peek();
        }else if(h1 == h2) {
            return ((double)minHeap.peek()+(double)maxHeap.peek())/2;
        }else {
            return h1>h2? (double)maxHeap.peek():(double)minHeap.peek();
        }
    }
};