480. Sliding Window Median

1class MedianFinder {
2    private PriorityQueue<Integer> minHeap = new PriorityQueue<>(); // holds all the larger elements
3    private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder()); // holds the smaller elements, reversed to act as a max-heap
4    private Map<Integer, Integer> countMap = new HashMap<>(); // keeps track of the number of instances a number has been delayed for removal
5    private int minHeapSize = 0;
6    private int maxHeapSize = 0;
7    private final int windowSize; // the size of the sliding window
8
9    public MedianFinder(int k) {
10        this.windowSize = k;
11    }
12
13    public void addNum(int num) {
14        // Add the number to the appropriate heap
15        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
16            maxHeap.offer(num);
17            maxHeapSize++;
18        } else {
19            minHeap.offer(num);
20            minHeapSize++;
21        }
22        rebalanceHeaps(); // Ensure both heaps remain in a valid state
23    }
24
25    public double findMedian() {
26        // If odd window size, the median is the top of the max heap, otherwise, the median is the average of the tops of both heaps
27        return (windowSize & 1) == 1 ? maxHeap.peek() : ((double) maxHeap.peek() + minHeap.peek()) / 2.0;
28    }
29
30    public void removeNum(int num) {
31        countMap.put(num, countMap.getOrDefault(num, 0) + 1); // Mark the number for delayed removal
32        if (num <= maxHeap.peek()) {
33            maxHeapSize--;
34            if (num == maxHeap.peek()) {
35                pruneHeap(maxHeap); // Remove elements that should be delayed from the maxHeap
36            }
37        } else {
38            minHeapSize--;
39            if (num == minHeap.peek()) {
40                pruneHeap(minHeap); // Remove elements that should be delayed from the minHeap
41            }
42        }
43        rebalanceHeaps(); // Ensure both heaps remain in a valid state
44    }
45
46    private void pruneHeap(PriorityQueue<Integer> heap) {
47        // Remove elements from the heap that were marked for delayed removal
48        while (!heap.isEmpty() && countMap.containsKey(heap.peek())) {
49            int heapTop = heap.peek();
50            countMap.put(heapTop, countMap.get(heapTop) - 1);
51
52            if (countMap.get(heapTop) == 0) {
53                countMap.remove(heapTop); // remove the entry from the map if the count goes to zero
54            }
55
56            heap.poll(); // remove the top element in the heap as it's already accounted for in the delayed removal
57        }
58    }
59
60    private void rebalanceHeaps() {
61        // Maintain the property that maxHeap has the same number of elements as minHeap, or 1 more
62        if (maxHeapSize > minHeapSize + 1) {
63            minHeap.offer(maxHeap.poll());
64            maxHeapSize--;
65            minHeapSize++;
66            pruneHeap(maxHeap); // Ensure maxHeap is pruned
67        } else if (maxHeapSize < minHeapSize) {
68            maxHeap.offer(minHeap.poll());
69            minHeapSize--;
70            maxHeapSize++;
71            pruneHeap(minHeap); // Ensure minHeap is pruned
72        }
73    }
74}
75
76class Solution {
77    public double[] medianSlidingWindow(int[] nums, int k) {
78        MedianFinder finder = new MedianFinder(k);
79        // Populate the MedianFinder with the initial sliding window
80        for (int i = 0; i < k; ++i) {
81            finder.addNum(nums[i]);
82        }
83
84        int n = nums.length;
85        double[] medians = new double[n - k + 1]; // Array to hold the medians
86        medians[0] = finder.findMedian(); // Find the median for the first window
87
88        // Move the sliding window and calculate the median for each window
89        for (int i = k; i < n; ++i) {
90            finder.addNum(nums[i]); // Add the new element to the sliding window
91            finder.removeNum(nums[i - k]); // Remove the oldest element from the sliding window
92            medians[i - k + 1] = finder.findMedian(); // Find the median for the current window
93        }
94        return medians; // Return the array of medians
95    }
96}
97

1#include <functional>
2#include <queue>
3#include <unordered_map>
4#include <vector>
5
6class MedianFinder {
7public:
8    // Constructor that initializes the finder with a window size 'k'.
9    MedianFinder(int k) : windowSize(k), smallSize(0), largeSize(0) {}
10
11    // Adds a number into the data structure.
12    void addNum(int num) {
13        if (small.empty() || num <= small.top()) {
14            small.push(num);
15            ++smallSize;
16        } else {
17            large.push(num);
18            ++largeSize;
19        }
20        // Ensure the heaps are balanced after the insertion.
21        rebalanceHeaps();
22    }
23
24    // Removes a number from the data structure.
25    void removeNum(int num) {
26        delayed[num]++;
27        if (num <= small.top()) {
28            --smallSize;
29            if (num == small.top()) {
30                prune(small);
31            }
32        } else {
33            --largeSize;
34            if (num == large.top()) {
35                prune(large);
36            }
37        }
38        // Ensure the heaps are balanced after the removal.
39        rebalanceHeaps();
40    }
41
42    // Returns the median of current numbers.
43    double findMedian() {
44        if (windowSize & 1) {
45            return small.top();
46        } else {
47            return (static_cast<double>(small.top()) + large.top()) * 0.5;
48        }
49    }
50
51private:
52    // Max heap to store the smaller half of the elements
53    std::priority_queue<int> small;
54    // Min heap to store the larger half of the elements
55    std::priority_queue<int, std::vector<int>, std::greater<int>> large;
56    // Hashmap to keep track of delayed removals when numbers are out of the current window
57    std::unordered_map<int, int> delayed;
58    // Sizes of the heaps
59    int smallSize, largeSize;
60    // Window size
61    int windowSize;
62
63    // Prunes/removes elements from the heap that should be delayed-removed.
64    template <typename T>
65    void prune(T& heap) {
66        while (!heap.empty() && delayed[heap.top()]) {
67            if (--delayed[heap.top()] == 0) {
68                delayed.erase(heap.top());
69            }
70            heap.pop();
71        }
72    }
73
74    // Balances the two heaps to maintain the median property.
75    void rebalanceHeaps() {
76        if (smallSize > largeSize + 1) {
77            large.push(small.top());
78            small.pop();
79            --smallSize;
80            ++largeSize;
81            // Re-prune in case the top element is meant to be delayed-removed.
82            prune(small);
83        } else if (smallSize < largeSize) {
84            small.push(large.top());
85            large.pop();
86            ++smallSize;
87            --largeSize;
88            // Re-prune in case the top element is meant to be delayed-removed.
89            prune(large);
90        }
91    }
92};
93
94class Solution {
95public:
96    // Calculates the median for each window of size 'k' in the array.
97    std::vector<double> medianSlidingWindow(std::vector<int>& nums, int k) {
98        MedianFinder finder(k);
99        // Initialize by adding the first 'k' elements.
100        for (int i = 0; i < k; ++i) {
101            finder.addNum(nums[i]);
102        }
103        // Store the first median.
104        std::vector<double> medians = {finder.findMedian()};
105        // Slide the window, one element at a time, maintaining the median.
106        for (int i = k; i < nums.size(); ++i) {
107            finder.addNum(nums[i]); // Add the new element to the window.
108            finder.removeNum(nums[i - k]); // Remove the element that is no longer in the window.
109            // Store the new median.
110            medians.push_back(finder.findMedian());
111        }
112        return medians;
113    }
114};
115

1// Create an alias for the comparison function type.
2type CompareFunction<T> = (a: T, b: T) => number;
3
4// Define PriorityQueue class template.
5class PriorityQueue<T> {
6    private data: T[] = [];
7    private compare: CompareFunction<T>;
8
9    constructor(compareFunction: CompareFunction<T>) {
10        this.compare = compareFunction;
11    }
12
13    private leftChildIndex(parentIndex: number): number {
14        return 2 * parentIndex + 1;
15    }
16
17    private rightChildIndex(parentIndex: number): number {
18        return 2 * parentIndex + 2;
19    }
20
21    private parentIndex(childIndex: number): number {
22        return Math.floor((childIndex - 1) / 2);
23    }
24
25    private hasLeftChild(index: number): boolean {
26        return this.leftChildIndex(index) < this.data.length;
27    }
28
29    private hasRightChild(index: number): boolean {
30        return this.rightChildIndex(index) < this.data.length;
31    }
32
33    private hasParent(index: number): boolean {
34        return this.parentIndex(index) >= 0;
35    }
36
37    private leftChild(index: number): T {
38        return this.data[this.leftChildIndex(index)];
39    }
40
41    private rightChild(index: number): T {
42        return this.data[this.rightChildIndex(index)];
43    }
44
45    private parent(index: number): T {
46        return this.data[this.parentIndex(index)];
47    }
48
49    private swap(indexOne: number, indexTwo: number): void {
50        let temp = this.data[indexOne];
51        this.data[indexOne] = this.data[indexTwo];
52        this.data[indexTwo] = temp;
53    }
54
55    private heapifyUp(): void {
56        let index = this.data.length - 1;
57        while (this.hasParent(index) && this.compare(this.parent(index), this.data[index]) > 0) {
58            this.swap(this.parentIndex(index), index);
59            index = this.parentIndex(index);
60        }
61    }
62
63    private heapifyDown(): void {
64        let index = 0;
65        while (this.hasLeftChild(index)) {
66            let smallerChildIndex = this.leftChildIndex(index);
67            if (this.hasRightChild(index) && this.compare(this.rightChild(index), this.leftChild(index)) < 0) {
68                smallerChildIndex = this.rightChildIndex(index);
69            }
70
71            if (this.compare(this.data[index], this.data[smallerChildIndex]) < 0) {
72                break;
73            } else {
74                this.swap(index, smallerChildIndex);
75            }
76            index = smallerChildIndex;
77        }
78    }
79
80    public isEmpty(): boolean {
81        return this.data.length === 0;
82    }
83
84    public peek(): T | undefined {
85        if (this.data.length === 0) return undefined;
86        return this.data[0];
87    }
88
89    public push(value: T): void {
90        this.data.push(value);
91        this.heapifyUp();
92    }
93
94    public pop(): T | undefined {
95        if (this.data.length === 0) return undefined;
96        if (this.data.length === 1) return this.data.pop();
97
98        const item = this.data[0];
99        this.data[0] = this.data.pop()!;
100        this.heapifyDown();
101        return item;
102    }
103
104    public size(): number {
105        return this.data.length;
106    }
107}
108
109// A max heap comparator for numbers.
110const maxHeapCompare: CompareFunction<number> = (a, b) => b - a;
111
112// A min heap comparator for numbers.
113const minHeapCompare: CompareFunction<number> = (a, b) => a - b;
114
115// Small values are stored in a max heap.
116const small: PriorityQueue<number> = new PriorityQueue(maxHeapCompare);
117// Large values are stored in a min heap.
118const large: PriorityQueue<number> = new PriorityQueue(minHeapCompare);
119
120// Map to keep track of values that need to be pruned from heaps.
121const delayed = new Map<number, number>();
122
123let windowSize = 0;
124let smallSize = 0;
125let largeSize = 0;
126
127function addNum(num: number): void {
128    if (small.isEmpty() || num <= small.peek()!) {
129        small.push(num);
130        smallSize++;
131    } else {
132        large.push(num);
133        largeSize++;
134    }
135    rebalanceHeaps();
136}
137
138function removeNum(num: number): void {
139    delayed.set(num, (delayed.get(num) || 0) + 1);
140    if (num <= small.peek()!) {
141        smallSize--;
142        if (num === small.peek()) {
143            prune(small);
144        }
145    } else {
146        largeSize--;
147        if (num === large.peek()) {
148            prune(large);
149        }
150    }
151    rebalanceHeaps();
152}
153
154function findMedian(): number {
155    if (windowSize % 2 === 1) {
156        return small.peek()!;
157    } else {
158        return (small.peek()! + large.peek()!) / 2.0;
159    }
160}
161
162function prune(heap: PriorityQueue<number>): void {
163    while (!heap.isEmpty() && delayed.get(heap.peek()!)! > 0) {
164        delayed.set(heap.peek()!, delayed.get(heap.peek()!)! - 1);
165        if (delayed.get(heap.peek()!) === 0) {
166            delayed.delete(heap.peek()!);
167        }
168        heap.pop();
169    }
170}
171
172function rebalanceHeaps(): void {
173    if (smallSize > largeSize + 1) {
174        large.push(small.pop()!);
175        smallSize--;
176        largeSize++;
177        prune(small);
178    } else if (smallSize < largeSize) {
179        small.push(large.pop()!);
180        smallSize++;
181        largeSize--;
182        prune(large);
183    }
184}
185
186function medianSlidingWindow(nums: number[], k: number): number[] {
187    windowSize = k;
188    const medians: number[] = [];
189
190    for (let i = 0; i < k; ++i) {
191        addNum(nums[i]);
192    }
193
194    medians.push(findMedian());
195
196    for (let i = k; i < nums.length; ++i) {
197        addNum(nums[i]);
198        removeNum(nums[i - k]);
199        medians.push(findMedian());
200    }
201
202    return medians;
203}
204