Leetcode: 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

Understand the problem:
Use two heaps. The maxHeap stores the number which is less than the current number. The minHeap stores the number which is greter than the current number. 
We also need to keep the two heaps balanced in size. 

For the method findMedian(), we need to check if the two heaps have the same size. If yes, there must be even number of elements so far, so the median is the average of the top of the minHeap and the maxHeap. If not, i.e. odd number of elements so far, the median is the top of the heap which one more element. 

Code (Java):

import java.util.*;

class MedianFinder {
    private PriorityQueue<Integer> leftPQ = 
        new PriorityQueue<>(Collections.reverseOrder());
    private PriorityQueue<Integer> rightPQ = new PriorityQueue<>();
    
    // Adds a number into the data structure.
    public void addNum(int num) {
        if (leftPQ.isEmpty() || num <= leftPQ.peek()) {
            leftPQ.offer(num);
        } else {
            rightPQ.offer(num);
        }
        
        // Rebalance the pqs
        if (leftPQ.size() - rightPQ.size() > 1) {
            rightPQ.offer(leftPQ.poll());
        } else if (rightPQ.size() - leftPQ.size() > 1) {
            leftPQ.offer(rightPQ.poll());
        }
    }

    // Returns the median of current data stream
    public double findMedian() {
        if (leftPQ.isEmpty() && rightPQ.isEmpty()) {
            throw new NoSuchElementException(); // if the queue is empty
        }
        
        if(leftPQ.isEmpty()) {
            return (double) rightPQ.peek();
        } else if (rightPQ.isEmpty()) {
            return (double) leftPQ.peek();
        } else if (leftPQ.size() > rightPQ.size()) {
            return (double) leftPQ.peek();
        } else if (rightPQ.size() > leftPQ.size()) {
            return (double) rightPQ.peek();
        } else {
            return (double) (leftPQ.peek() + rightPQ.peek()) / 2;
        }
    }
};

// Your MedianFinder object will be instantiated and called as such:
// MedianFinder mf = new MedianFinder();
// mf.addNum(1);
// mf.findMedian();

Update 3/19/19:

public class Solution {
    /**
     * @param nums: A list of integers
     * @return: the median of numbers
     */
    public int[] medianII(int[] nums) {
        if (nums == null || nums.length == 0) {
            return new int[0];
        }

        PriorityQueue<Integer> maxPQ = new PriorityQueue<>(Collections.reverseOrder());
        PriorityQueue<Integer> minPQ = new PriorityQueue<>();

        int[] ans = new int[nums.length];

        for (int i = 0; i < nums.length; i++) {
            int num = nums[i];
            if (maxPQ.isEmpty() || num <= maxPQ.peek()) {
                maxPQ.offer(num);
            } else {
                minPQ.offer(num);
            }

            rebalacne(maxPQ, minPQ);

            ans[i] = maxPQ.peek();
        }
        
        return ans;
    }

    private void rebalacne(PriorityQueue<Integer> maxPQ, PriorityQueue<Integer> minPQ) {        
        if (maxPQ.size() - minPQ.size() > 1) {
            minPQ.offer(maxPQ.poll());
        } else if (minPQ.size() - maxPQ.size() > 0) {
            maxPQ.offer(minPQ.poll());
        }
    }
}