[刷题]Data Stream Median

[LintCode]Data Stream Median

public class Solution {
    /**
     * @param nums: A list of integers.
     * @return: the median of numbers
     */
    public int[] medianII(int[] nums) {
        // 2015-09-04
        // median 中位数
        // 左堆是最大堆，右堆是最小堆
        // 只有两种情况：两堆元素个数相等，或左堆多一个
        if (nums == null || nums.length == 0) {
            return new int[0];
        }
        
        int initialSize = nums.length;
        int[] medianList = new int[initialSize];
        // 如果不用initialSize会报错，不知道为什么
        PriorityQueue<Integer> leftHeap = new PriorityQueue<Integer> (initialSize,  
        new Comparator<Integer> () {  
            public int compare(Integer a, Integer b) {  
                return b - a;  
            }  
        }); 
        PriorityQueue<Integer> rightHeap = new PriorityQueue<>();
        
        for (int i = 0; i < nums.length; i++) {
            // 两堆都为空
            if (leftHeap.isEmpty()) {
                leftHeap.offer(nums[i]);
                medianList[i] = leftHeap.peek();
                continue;
            }
            
            if (nums[i] < leftHeap.peek()) {
                // 这个元素得放在左队
                if (leftHeap.size() == rightHeap.size()) {
                    leftHeap.offer(nums[i]);
                } else {
                    leftHeap.offer(nums[i]);
                    rightHeap.offer(leftHeap.poll());
                }
            } else {
                // 这个元素得放在右堆
                if (leftHeap.size() == rightHeap.size()) {
                    rightHeap.offer(nums[i]);
                    leftHeap.offer(rightHeap.poll());
                } else {
                    rightHeap.offer(nums[i]);
                }
            }
            medianList[i] = leftHeap.peek();
        }
        
        return medianList;
    }
}