GuilinDev

Lc0480

05 August 2008

480. Sliding Window Median

滑动窗口求中位数

Slide through each of the window of size K and calculate median of each window.

Time Complexity : O(nlogn)

Space Complexity: O(n)

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
class Solution {

    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    
    private void removeFromHeap(int n) {
        if (n <= maxHeap.peek()) {
            if (maxHeap.size() > minHeap.size()) {
                maxHeap.remove(n);
            } else {
                maxHeap.remove(n);
                maxHeap.add(minHeap.poll());
            }
        } else {
            if (maxHeap.size() > minHeap.size()) {
                minHeap.remove(n);
                minHeap.add(maxHeap.poll());
            } else {
                minHeap.remove(n);
            }
        }
    }

    private void insertToHeap(int num) {
        if (maxHeap.isEmpty()) {
            maxHeap.add(num);
        } else {
            if (maxHeap.size() > minHeap.size()) {
                if (maxHeap.peek() < num) {
                    minHeap.add(num);
                } else {
                    minHeap.add(maxHeap.poll());
                    maxHeap.add(num);
                }
            } else {
                if (maxHeap.peek() <= num) {
                    minHeap.add(num);
                    maxHeap.add(minHeap.poll());
                } else {
                    maxHeap.add(num);
                }
            }
        }
    }

    private double median() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        }
        long sum = (long) maxHeap.peek() + (long) minHeap.peek();
        return (sum) / 2.0;
    }

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

        for (int i = 0; i <= nums.length; ++i) {
            if (i >= k) {
                res[i - k] = median();
                removeFromHeap(nums[i - k]);
            }

            if (i < nums.length) {
                insertToHeap(nums[i]);
            }
        }

        return res;
    }
}