[LintCode/LeetCode] Find Median From / Data Stream Median

Problem

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

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.

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

Total run time in O(nlogn).

Tags

LintCode Copyright Heap Priority Queue Google

Note

建立两个堆，一个堆就是PriorityQueue本身，也就是一个最小堆；另一个要写一个Comparator，使之成为一个最大堆。我们把遍历过的数组元素对半分到两个堆里，更大的数放在最小堆，较小的数放在最大堆。为什么这么分呢？因为要从maxHeap堆顶取较小的一半元素中最大的那个，而对另一半较大的数，我们并不关心。
同时，确保最大堆的size比最小堆大1，才能从最大堆的顶端返回median。

Solution

public class Solution {
    public int[] medianII(int[] nums) {
        if (nums == null || nums.length == 0) return new int[0];
        int[] res = new int[nums.length];
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(16, new Comparator<Integer>() {
            public int compare(Integer x, Integer y) {
                return y-x;
            }
        });
        res[0] = nums[0];
        maxHeap.add(nums[0]);
        for (int i = 1; i < nums.length; i++) {
            int max = maxHeap.peek();
            if (nums[i] <= max) maxHeap.add(nums[i]);
            else minHeap.add(nums[i]);
            if (maxHeap.size() > minHeap.size()+1) minHeap.add(maxHeap.poll());
            else if (maxHeap.size() < minHeap.size()) maxHeap.add(minHeap.poll());
            res[i] = maxHeap.peek();
        }
        return res;
    }
}

LeetCode Version

public class MedianFinder {
    PriorityQueue<Integer> minheap = new PriorityQueue<>();
    PriorityQueue<Integer> maxheap = new PriorityQueue<>(1, new Comparator<Integer>() {
        public int compare(Integer i1, Integer i2) {
            return i2-i1;
        }
    });
    // Or we can use: = new PriorityQueue<>(1, Collections.reverseOrder());
    
    // Adds a number into the data structure.
    public void addNum(int num) {
        maxheap.offer(num);
        minheap.offer(maxheap.poll());
        if (maxheap.size() < minheap.size()) maxheap.offer(minheap.poll());
        else if (maxheap.size() > minheap.size()) minheap.offer(maxheap.poll());
        //or (maxheap.size() > minheap.size()+1)
    }

    // Returns the median of current data stream
    public double findMedian() {
        if(maxheap.size() == minheap.size()) return (maxheap.peek()+minheap.peek())/2.0;
        return maxheap.peek();
    }
};