LC480.滑动窗口中位数（双堆+滑动窗口）

问题

题解

• 题目同数据流的中位数一样，都是使用双堆维护，详细可以看那一题。
• 不同点在于，给数据流加了一个滑动窗口，数据需要从窗口中移除，因此只需要判断要被移除的元素是在哪个堆中即可，找到后删除，heap删除元素的时间复杂度为O(n)

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>((Collections.reverseOrder()));
        PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
        
        double[] res = new double[nums.length - k + 1];
        for(int i = 0; i < nums.length; i ++){
            if(maxHeap.isEmpty() || nums[i] <= maxHeap.peek()){
                maxHeap.add(nums[i]);
            }else{
                minHeap.add(nums[i]);
            }
            balanceHeap(maxHeap, minHeap);

            if(i - k + 1 >= 0){
                if(k % 2 == 1){//如果元素个数为奇数
                    res[i - k + 1] = maxHeap.peek();
                }else{//元素个数为偶数
                    res[i - k + 1] = maxHeap.peek() / 2.0 + minHeap.peek() / 2.0;
                }
            	//移除元素，需要先判断元素在哪个堆中
                int elementToRemove = nums[i - k + 1];
                if(elementToRemove <= maxHeap.peek()){
                    maxHeap.remove(elementToRemove);
                }else{
                    minHeap.remove(elementToRemove);
                }
                balanceHeap(maxHeap, minHeap);
            }
        }

        return res;
    }
    public void balanceHeap(PriorityQueue<Integer> maxHeap, PriorityQueue<Integer> minHeap){
        if(maxHeap.size() < minHeap.size()){
            maxHeap.add(minHeap.poll());
        }else if(maxHeap.size() > minHeap.size() + 1){
            minHeap.add(maxHeap.poll());
        }
    }
}