480. Sliding Window Median 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
class Solution {
    private static class DualHeap {
        private final Queue<Integer> maxHeap;
        private final Queue<Integer> minHeap;
        private final int k;
        private int maxHeapSize;
        private int minHeapSize;
        private Map<Integer, Integer> deletedValueToCountMap;

        public DualHeap(int k) {
            this.k = k;
            this.maxHeap = new PriorityQueue<>(Comparator.reverseOrder()); // 注意此处需要使用比较器而不能使用减号，因为可能越界
            this.minHeap = new PriorityQueue<>();
            this.deletedValueToCountMap = new HashMap<>();
        }

        public double getMedian() {
            if ((k & 1) == 1) {
                return maxHeap.peek();
            } else {
                // 注意 double 转换必须放在前面，以规避越界问题，可使用以下入参测试：nums = [2147483647, 2147483647], k = 2
                // 会越界的写法：(maxHeap.peek() + minHeap.peek()) / 2.0;
                return ((double) maxHeap.peek() + minHeap.peek()) / 2;
            }
        }

        // TODO 该方法实现可优化
        public void add(int num) {
            if (maxHeapSize == minHeapSize) {
                // 往 maxHeap 中添加
                if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
                    maxHeap.offer(num);
                } else {
                    minHeap.offer(num);
                    maxHeap.offer(minHeap.poll());
                    prune(minHeap); // 堆顶元素变化的堆都需要尝试 prune, 不能让已删除元素停留在堆顶
                }
                maxHeapSize++;
            } else {
                // 往 minHeap 中添加
                if (minHeap.isEmpty() || num >= minHeap.peek()) {
                    minHeap.offer(num);
                } else {
                    maxHeap.offer(num);
                    minHeap.offer(maxHeap.poll());
                    prune(maxHeap); // 堆顶元素变化的堆都需要尝试 prune, 不能让已删除元素停留在堆顶
                }
                minHeapSize++;
            }
        }

        public void remove(int num) {
            deletedValueToCountMap.put(num, deletedValueToCountMap.getOrDefault(num, 0) + 1);
            if (num <= maxHeap.peek()) {
                maxHeapSize--;
                if (num == maxHeap.peek()) {
                    prune(maxHeap);
                }
            } else {
                minHeapSize--;
                if (num == minHeap.peek()) {
                    prune(minHeap);
                }
            }

            // 由于删除了堆中的元素，所以此时堆顶元素可能不满足中位数的性质，所以此时需要校验 maxHeapSize 与 minHeapSize, 检查 maxHeapSize - minHeapSize 是否等于 0 或 1
            makeBalance();
        }

        private void makeBalance() {
            if (maxHeapSize < minHeapSize) {
                maxHeap.offer(minHeap.poll());
                maxHeapSize++;
                minHeapSize--;
                prune(minHeap); // 堆顶元素变化的堆都需要尝试 prune, 不能让已删除元素停留在堆顶
            } else if (maxHeapSize - minHeapSize > 1) {
                minHeap.offer(maxHeap.poll());
                minHeapSize++;
                maxHeapSize--;
                prune(maxHeap); // 堆顶元素变化的堆都需要尝试 prune, 不能让已删除元素停留在堆顶
            }
        }

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

    }

    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] res = new double[nums.length - k + 1];

        DualHeap dualHeap = new DualHeap(k);
        for (int i = 0; i < k; i++) {
            dualHeap.add(nums[i]);
        }
        res[0] = dualHeap.getMedian();
        int resIndex = 1;

        for (int i = k; i < nums.length; i++) {
            dualHeap.add(nums[i]);
            dualHeap.remove(nums[i - k]);
            res[resIndex++] = dualHeap.getMedian();
        }

        return res;
    }
}

References

480. Sliding Window Median
The biggest integer value can be represented by double in Java