[leetCode]480. 滑动窗口中位数

题目

https://leetcode-cn.com/problems/sliding-window-median/

双优先队列加延迟删除

由于要得到一个窗口内的中位数，可以使用两个优先队列small（大顶堆）、large（小顶堆），一个维护窗口内元数排序后较小的一半，一个维护较大的一半，如果元素个数为奇数那个中位数为small的堆顶元素，如果元素个数为偶数那么中位数就是small与large的堆顶元素除以2，此题的难点在于怎样保持两个堆的元素数量各占一半。
一开始将元素加入small中，如果small不为空则后续元素与small堆顶元素比较，如果小于small的堆顶元素则加入small中，如果大于small的堆顶元素则加入到large中。为了维持small与large的元素数量关系，当small的元素数量比large多2个时，就将small的堆顶元素加入到large中，如果small的元素数量比large少一个就将large的堆顶元素加入到small中。

由于滑动窗口需要删除元素，而堆不能删除中间位置的元素，因此采用延迟删除的方法，将要删除的元素加入哈希表中，如果该元素到堆顶时再把该元素删除。因此使用辅助函数prune(heap)来不断删除堆顶元素，保证了堆顶元素一定不是要删除的元素，这样就方便计算中位数了。

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        DualHeap dh = new DualHeap(k);
        for (int i = 0; i < k; i++) {
            dh.insert(nums[i]);
        }
        double[] ans = new double[nums.length - k + 1];
        ans[0] = dh.getMediean();
        for (int i = k; i < nums.length; i++) {
            dh.insert(nums[i]);
            dh.erase(nums[i - k]);
            ans[i - k + 1] = dh.getMediean();
        }
        return ans;
    }

    class DualHeap {
        // 大根堆，维护较小的一般元素
        private PriorityQueue<Integer> small;
        // 小根堆，维护较大的一般元素
        private PriorityQueue<Integer> large;
        // 哈希表，记录被延迟删除的元素 key为要删除的元素，value为需要删除的次数
        private Map<Integer, Integer> delayed;
        // 滑动窗口大小
        private int k;
        // 由于延迟删除的操作，所以需要两个变量维护 small 和 large 的实际大小
        private int smallSize, largeSize;

        public DualHeap(int k) {
            this.k = k;
            small = new PriorityQueue<Integer>((o1, o2) -> o2.compareTo(o1));
            large = new PriorityQueue<Integer>();
            delayed = new HashMap<>();
            smallSize = 0;
            largeSize = 0;
        }

        public double getMediean() {
            return (k & 1) == 1 ? small.peek() : ((double)small.peek() + large.peek()) / 2;
        }

        public void insert(int num) {
            if (small.isEmpty() || num < small.peek()) {
                small.offer(num);
                smallSize++;
            } else {
                large.offer(num);
                largeSize++;
            }
            makeBalance();
        }

        public void erase(int num) {
            delayed.put(num, delayed.getOrDefault(num, 0) + 1);
            if (num <= small.peek()) {
                --smallSize;
                if (num == small.peek()) {
                    prune(small);
                } 
            } else {
                --largeSize;
                if (num == large.peek()) { 
                    prune(large);
                }
            }
            makeBalance();
        }


        private void prune(PriorityQueue<Integer> heap) {
            while (!heap.isEmpty()) {
                int num = heap.peek();
                if (delayed.containsKey(num)) {
                    delayed.put(num, delayed.get(num) - 1);
                    if (delayed.get(num) == 0) {
                        delayed.remove(num);
                    }
                    heap.poll();
                } else {
                    break;
                }
            }
        }

        private void makeBalance() {
            if (smallSize > largeSize + 1) {
                large.offer(small.poll());
                largeSize++;
                smallSize--;
                prune(small);
            } else if (smallSize < largeSize) {
                small.offer(large.poll());
                largeSize--;
                smallSize++;
                prune(large);
            }
        }
    }
}