LintCode-Median II

Numbers keep coming, return the median of numbers at every time a new number added.

Example

For numbers coming list: [1, 2, 3, 4, 5], return [1, 1, 2, 2, 3]

For numbers coming list: [4, 5, 1, 3, 2, 6, 0], return [4, 4, 4, 3, 3, 3, 3]

For numbers coming list: [2, 20, 100], return [2, 2, 20]

Challenge

O(nlogn) time

Clarification

What's the definition of Median?

• Median is the number that in the middle of a sorted array. If there are n numbers in a sorted array A, the median is A[(n-1)/2].
• For example, if A=[1,2,3], median is 2. If A=[1,19], median is 1

Analysis:

We maintain a max heap and a min heap. At any index i, the max heap stores the elements small or equal than the current median and the min heap stores the elements that are larger than the current median. Because in the problem, we select the median as at the position (n-1)/2, so we should always keep the size of max heap equal or one larger than the min heap. At odd index, we try to rebalance the size of max heap and min heap, while at even index, we try to make the size of max heap one larger than that of the min heap. By doing this, at each position, after inserting A[i] into the heaps, the root value of the max heap is the median.

NOTE: if we select (n-1)/2+1 as median, we can then keep the min heap equal or one larger than the max heap, the root value of the min heap is the median.

Solution:

I implement a Heap class with only insert and popHeapRoot functions.

class Heap{
    private int[] nodes;
    private int size;
    private boolean isMaxHeap;
    
    public Heap(int capa, boolean isMax){
        nodes = new int[capa];
        size = 0;
        isMaxHeap = isMax;
    }
    
    public boolean isEmpty(){
        if (size==0) return true;
        else return false;
    }

    public int getHeapRootValue(){
        //should throw exception when size==0;
        return nodes[0];
    }

    private void swap(int x, int y){
        int temp = nodes[x];
        nodes[x] = nodes[y];
        nodes[y] = temp;
    }

    public boolean insert(int val){
        if (size==nodes.length) return false;
        size++;
        nodes[size-1]=val;
        //check its father iteratively.
        int cur = size-1;
        int father = (cur-1)/2;
        while (father>=0 && ((isMaxHeap && nodes[cur]>nodes[father]) || (!isMaxHeap && nodes[cur]<nodes[father]))){
            swap(cur,father);
            cur = father;
            father = (cur-1)/2;
        }
        return true;
    }

    private void shiftDown(int ind){
        int left = (ind+1)*2-1;
        int right = (ind+1)*2;
        while (left<size || right<size){
            if (isMaxHeap){
                int leftVal = (left<size) ? nodes[left] : Integer.MIN_VALUE;
                int rightVal = (right<size) ? nodes[right] : Integer.MIN_VALUE;
                int next = (leftVal>=rightVal) ? left : right;
                if (nodes[ind]>nodes[next]) break;
                else {
                    swap(ind,next);
                    ind = next;
                    left = (ind+1)*2-1;
                    right = (ind+1)*2;
                }
            } else {
                int leftVal = (left<size) ? nodes[left] : Integer.MAX_VALUE;
                int rightVal = (right<size) ? nodes[right] : Integer.MAX_VALUE;
                int next = (leftVal<=rightVal) ? left : right;
                if (nodes[ind]<nodes[next]) break;
                else {
                    swap(ind,next);
                    ind = next;
                    left = (ind+1)*2-1;
                    right = (ind+1)*2;
                }
            }
        }
    }    

    public int popHeapRoot(){
        //should throw exception, when heap is empty.

        int rootVal = nodes[0];
        swap(0,size-1);
        size--;
        if (size>0) shiftDown(0);
        return rootVal;
    }
}
        
            
    

public class Solution {
    /**
     * @param nums: A list of integers.
     * @return: the median of numbers
     */
    public int[] medianII(int[] nums) {
        int[] res = new int[nums.length];
        Heap maxHeap = new Heap(nums.length,true);
        Heap minHeap = new Heap(nums.length,false);
        maxHeap.insert(nums[0]);
        res[0] = nums[0];
        for (int i=1;i<nums.length;i++)
            if (i %2 == 1) {    //i is odd index.
                int median = maxHeap.getHeapRootValue();
                if (nums[i]<median){
                    maxHeap.popHeapRoot();
                    minHeap.insert(median);
                    maxHeap.insert(nums[i]);
                    res[i] = maxHeap.getHeapRootValue();
                } else {
                    res[i] = median;
                    minHeap.insert(nums[i]);
                }
            } else {        //i is even index.
                int median = maxHeap.getHeapRootValue();
                if (nums[i]<median){
                    maxHeap.insert(nums[i]);
                } else {
                    minHeap.insert(nums[i]);
                    int val = minHeap.popHeapRoot();
                    maxHeap.insert(val);
                }
                res[i] = maxHeap.getHeapRootValue();
            }

        return res;
    }
}