Facebook Pixel

480. Sliding Window Median

HardArrayHash TableSliding WindowHeap (Priority Queue)
Leetcode Link

Problem Description

You are given an integer array nums and an integer k. You need to find the median of every contiguous subarray of size k using a sliding window approach.

The median is defined as:

  • For an odd-sized list: the middle element when sorted
  • For an even-sized list: the average of the two middle elements when sorted

For example:

  • If arr = [2, 3, 4], the median is 3 (the middle element)
  • If arr = [1, 2, 3, 4], the median is (2 + 3) / 2 = 2.5 (average of two middle elements)

The sliding window works as follows:

  • Start with a window of size k at the leftmost position of the array
  • Calculate the median of the current window
  • Move the window one position to the right (remove the leftmost element, add the next element from the right)
  • Continue until the window reaches the rightmost position

Your task is to return an array containing the median of each window position as the window slides from left to right across the array. The answers will be accepted if they are within 10^-5 of the actual value.

For instance, if nums = [1, 3, -1, -3, 5, 3, 6, 7] and k = 3:

  • Window [1, 3, -1] → median = 1
  • Window [3, -1, -3] → median = -1
  • Window [-1, -3, 5] → median = -1
  • And so on...

The result would be an array of all these medians.

Quick Interview Experience
Help others by sharing your interview experience
Have you seen this problem before?

Intuition

Finding the median requires knowing the middle element(s) of a sorted sequence. If we were to sort the window every time it slides, we'd need O(k log k) time for each position, which is inefficient.

The key insight is that we don't need the entire sorted order - we only need to efficiently access the middle element(s). This naturally leads us to think about dividing our elements into two halves: the smaller half and the larger half.

If we maintain these two halves such that:

  • The smaller half contains the smallest ⌈k/2⌉ elements
  • The larger half contains the remaining elements
  • We can quickly access the maximum of the smaller half and the minimum of the larger half

Then the median would be:

  • For odd k: the maximum of the smaller half
  • For even k: the average of the maximum of the smaller half and the minimum of the larger half

This structure points us toward using two heaps:

  • A max-heap for the smaller half (to quickly get the largest element in this half)
  • A min-heap for the larger half (to quickly get the smallest element in this half)

However, the sliding window introduces a challenge: we need to remove elements that fall out of the window. Removing arbitrary elements from a heap is expensive (O(k) time).

Here's where lazy deletion comes in. Instead of immediately removing an element from the heap, we:

  1. Mark it for deletion in a separate dictionary
  2. Only actually remove it when it appears at the top of a heap
  3. Keep track of the "logical" sizes of our heaps (actual size minus elements marked for deletion)

This way, elements that are buried in the heap and marked for deletion don't affect our median calculation (since they're not at the top), and we only pay the cost of removal when absolutely necessary.

The rebalancing ensures that our two halves maintain the correct size relationship, so the median can always be computed from the tops of the heaps.

Solution Approach

The solution implements a MedianFinder class that maintains two heaps with lazy deletion to efficiently track the median of a sliding window.

Data Structures Used

  1. Two Heaps:

    • small: A max-heap (implemented using negative values in Python's min-heap) storing the smaller half of elements
    • large: A min-heap storing the larger half of elements

  2. Lazy Deletion Dictionary:

    • delayed: A dictionary tracking elements to be removed and their counts
    • small_size and large_size: Logical sizes of the heaps (excluding delayed elements)

Key Methods Implementation

1. add_num(num):

  • If small is empty or num ≤ -small[0] (max of smaller half), add to small
  • Otherwise, add to large
  • Increment the appropriate size counter
  • Call rebalance() to maintain heap size invariant

2. remove_num(num):

  • Mark num for lazy deletion by incrementing delayed[num]
  • Determine which heap logically contains num by comparing with -small[0]
  • Decrement the appropriate size counter
  • If num is at the top of its heap, call prune() to clean it up immediately

3. prune(pq):

  • While the top element of heap pq is in the delayed dictionary:
    • Decrement its count in delayed
    • If count reaches 0, remove from delayed
    • Pop the element from the heap
  • This ensures the heap's top is always a valid (non-delayed) element

4. rebalance():

  • Maintains the invariant: small_size should be either equal to large_size or large_size + 1
  • If small_size > large_size + 1: Move top of small to large
  • If small_size < large_size: Move top of large to small
  • After moving elements, prune the source heap to ensure its new top is valid

5. find_median():

  • If k is odd: Return -small[0] (the max of smaller half)
  • If k is even: Return (-small[0] + large[0]) / 2 (average of two middle elements)

Main Algorithm Flow

  1. Initialize MedianFinder with window size k
  2. Add the first k elements to build the initial window
  3. Calculate the first median
  4. For each remaining element in nums:
    • Add the new element entering the window
    • Remove the element leaving the window (lazy deletion)
    • Calculate and store the new median

Time Complexity Analysis

  • Adding an element: O(log k) for heap insertion
  • Removing an element: O(1) for marking deletion, O(log k) amortized for actual removal
  • Finding median: O(1) as we just access heap tops
  • Overall: O(n log k) where n is the array length

The lazy deletion optimization is crucial here - it prevents us from having to search through the heap to remove elements, which would be O(k) per removal.

Ready to land your dream job?

Unlock your dream job with a 3-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through the solution with nums = [1, 3, -1, -3, 5, 3] and k = 3.

Initial Window: [1, 3, -1]

  1. Start with empty heaps: small = [], large = []
  2. Add 1: Since small is empty, add to small.
    • small = [-1] (using negative for max-heap), small_size = 1
  3. Add 3: Compare with max of small (-(-1) = 1). Since 3 > 1, add to large.
    • large = [3], large_size = 1
  4. Add -1: Compare with max of small (1). Since -1 ≤ 1, add to small.
    • small = [-1, 1] (heap form), small_size = 2
  5. Rebalance: small_size (2) > large_size (1) + 1, so move top of small to large.
    • Move 1 from small to large
    • small = [1] (contains -1), small_size = 1
    • large = [1, 3] (heap form), large_size = 2
  6. Rebalance again: small_size (1) < large_size (2), so move top of large to small.
    • Move 1 from large to small
    • small = [1, -1] (contains -1, 1), small_size = 2
    • large = [3], large_size = 1
  7. Find median: k is odd, so median = -small[0] = -(-1) = 1

Slide to [3, -1, -3]

  1. Remove 1 (leaving element): Mark in delayed = {1: 1}. Since 1 is at top of small, prune it immediately.
    • After pruning: small = [1] (contains -1), small_size = 1
  2. Add -3 (entering element): Compare with max of small (-1). Since -3 ≤ -1, add to small.
    • small = [1, 3] (contains -1, -3), small_size = 2
  3. Rebalance: small_size (2) > large_size (1) + 1, so move top of small to large.
    • Move -1 from small to large
    • small = [3] (contains -3), small_size = 1
    • large = [-1, 3] (heap form), large_size = 2
  4. Rebalance again: small_size (1) < large_size (2), so move top of large to small.
    • Move -1 from large to small
    • small = [1, 3] (contains -3, -1), small_size = 2
    • large = [3], large_size = 1
  5. Find median: k is odd, so median = -small[0] = -1 = -1

Slide to [-1, -3, 5]

  1. Remove 3 (leaving element): Mark in delayed = {3: 1}. Since 3 is in large and at the top, prune it.
    • After pruning: large = [], large_size = 0
  2. Add 5 (entering element): Compare with max of small (-1). Since 5 > -1, add to large.
    • large = [5], large_size = 1
  3. No rebalancing needed (small_size = 2, large_size = 1 is valid)
  4. Find median: k is odd, so median = -small[0] = -1 = -1

Slide to [-3, 5, 3]

  1. Remove -1 (leaving element): Mark in delayed = {-1: 1}. Since -1 is at top of small, prune it.
    • After pruning: small = [3] (contains -3), small_size = 1
  2. Add 3 (entering element): Compare with max of small (-3). Since 3 > -3, add to large.
    • large = [3, 5] (heap form), large_size = 2
  3. Rebalance: small_size (1) < large_size (2), so move top of large to small.
    • Move 3 from large to small
    • small = [3, -3] (contains -3, 3), small_size = 2
    • large = [5], large_size = 1
  4. Find median: k is odd, so median = -small[0] = -(-3) = 3

Final result: [1, -1, -1, 3]

The key insight is how lazy deletion allows us to mark elements for removal without searching through the heaps, and how the two-heap structure with rebalancing maintains quick access to the median at all times.

Solution Implementation

1from collections import defaultdict
2from heapq import heappush, heappop
3from typing import List
4
5class MedianFinder:
6    """
7    A data structure that maintains the median of a sliding window of size k.
8    Uses two heaps (max heap for smaller half, min heap for larger half) with lazy deletion.
9    """
10  
11    def __init__(self, k: int):
12        """
13        Initialize the MedianFinder with window size k.
14      
15        Args:
16            k: The size of the sliding window
17        """
18        self.k = k
19        # Max heap for smaller half (negated values for max heap behavior)
20        self.small = []
21        # Min heap for larger half
22        self.large = []
23        # Track elements to be lazily deleted with their counts
24        self.delayed = defaultdict(int)
25        # Actual size of elements in small heap (excluding delayed deletions)
26        self.small_size = 0
27        # Actual size of elements in large heap (excluding delayed deletions)
28        self.large_size = 0
29
30    def add_num(self, num: int) -> None:
31        """
32        Add a number to the data structure.
33      
34        Args:
35            num: The number to add
36        """
37        # Determine which heap to add the number to
38        if not self.small or num <= -self.small[0]:
39            # Add to small heap (negate for max heap behavior)
40            heappush(self.small, -num)
41            self.small_size += 1
42        else:
43            # Add to large heap
44            heappush(self.large, num)
45            self.large_size += 1
46      
47        # Ensure heaps remain balanced
48        self.rebalance()
49
50    def find_median(self) -> float:
51        """
52        Find the median of current elements in the window.
53      
54        Returns:
55            The median value as a float
56        """
57        # If k is odd, median is the top of small heap
58        # If k is even, median is average of tops of both heaps
59        return -self.small[0] if self.k & 1 else (-self.small[0] + self.large[0]) / 2
60
61    def remove_num(self, num: int) -> None:
62        """
63        Remove a number from the data structure using lazy deletion.
64      
65        Args:
66            num: The number to remove
67        """
68        # Mark the number for lazy deletion
69        self.delayed[num] += 1
70      
71        # Update size and prune if the number is at the top
72        if num <= -self.small[0]:
73            self.small_size -= 1
74            if num == -self.small[0]:
75                self.prune(self.small)
76        else:
77            self.large_size -= 1
78            if num == self.large[0]:
79                self.prune(self.large)
80      
81        # Ensure heaps remain balanced
82        self.rebalance()
83
84    def prune(self, heap: List[int]) -> None:
85        """
86        Remove all delayed elements from the top of the heap.
87      
88        Args:
89            heap: The heap to prune (either self.small or self.large)
90        """
91        # Determine sign for value conversion (-1 for small heap, 1 for large heap)
92        sign = -1 if heap is self.small else 1
93      
94        # Remove all delayed elements from the top
95        while heap and sign * heap[0] in self.delayed:
96            # Decrease the delay count
97            self.delayed[sign * heap[0]] -= 1
98            # Remove from delayed dict if count reaches 0
99            if self.delayed[sign * heap[0]] == 0:
100                self.delayed.pop(sign * heap[0])
101            # Remove from heap
102            heappop(heap)
103
104    def rebalance(self) -> None:
105        """
106        Rebalance the two heaps to maintain the median property.
107        Small heap should have equal or one more element than large heap.
108        """
109        # If small heap has too many elements, move one to large heap
110        if self.small_size > self.large_size + 1:
111            heappush(self.large, -heappop(self.small))
112            self.small_size -= 1
113            self.large_size += 1
114            # Clean up any delayed elements exposed at the top
115            self.prune(self.small)
116        # If large heap has more elements, move one to small heap
117        elif self.small_size < self.large_size:
118            heappush(self.small, -heappop(self.large))
119            self.large_size -= 1
120            self.small_size += 1
121            # Clean up any delayed elements exposed at the top
122            self.prune(self.large)
123
124
125class Solution:
126    """
127    Solution for the Sliding Window Median problem.
128    """
129  
130    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
131        """
132        Find the median of each sliding window of size k in the array.
133      
134        Args:
135            nums: The input array of numbers
136            k: The size of the sliding window
137          
138        Returns:
139            List of medians for each window position
140        """
141        # Initialize the median finder with window size k
142        finder = MedianFinder(k)
143      
144        # Add the first k elements to initialize the window
145        for num in nums[:k]:
146            finder.add_num(num)
147      
148        # Calculate the median for the first window
149        result = [finder.find_median()]
150      
151        # Slide the window through the rest of the array
152        for i in range(k, len(nums)):
153            # Add new element entering the window
154            finder.add_num(nums[i])
155            # Remove element leaving the window
156            finder.remove_num(nums[i - k])
157            # Calculate and store the median for current window
158            result.append(finder.find_median())
159      
160        return result
161
1/**
2 * A data structure that maintains the median of a sliding window of size k.
3 * Uses two heaps (max heap for smaller half, min heap for larger half) to efficiently track the median.
4 * Implements lazy deletion to handle element removal efficiently.
5 */
6class MedianFinder {
7    // Max heap to store the smaller half of numbers
8    private PriorityQueue<Integer> smallerHalf = new PriorityQueue<>(Comparator.reverseOrder());
9  
10    // Min heap to store the larger half of numbers
11    private PriorityQueue<Integer> largerHalf = new PriorityQueue<>();
12  
13    // Map to track numbers that need to be removed (lazy deletion)
14    private Map<Integer, Integer> delayedRemovals = new HashMap<>();
15  
16    // Actual size of smaller half (excluding delayed removals)
17    private int smallerHalfSize;
18  
19    // Actual size of larger half (excluding delayed removals)
20    private int largerHalfSize;
21  
22    // Window size
23    private int windowSize;
24
25    /**
26     * Constructor to initialize the MedianFinder with window size k
27     * @param k the size of the sliding window
28     */
29    public MedianFinder(int k) {
30        this.windowSize = k;
31    }
32
33    /**
34     * Adds a number to the data structure
35     * @param num the number to add
36     */
37    public void addNum(int num) {
38        // Determine which heap to add the number to
39        if (smallerHalf.isEmpty() || num <= smallerHalf.peek()) {
40            smallerHalf.offer(num);
41            ++smallerHalfSize;
42        } else {
43            largerHalf.offer(num);
44            ++largerHalfSize;
45        }
46      
47        // Maintain heap balance after insertion
48        rebalance();
49    }
50
51    /**
52     * Finds and returns the median of current window
53     * @return the median value
54     */
55    public double findMedian() {
56        // If window size is odd, return the top of smaller half
57        // If window size is even, return average of tops of both heaps
58        return (windowSize & 1) == 1 
59            ? smallerHalf.peek() 
60            : ((double) smallerHalf.peek() + largerHalf.peek()) / 2;
61    }
62
63    /**
64     * Removes a number from the data structure (lazy deletion)
65     * @param num the number to remove
66     */
67    public void removeNum(int num) {
68        // Mark the number for delayed removal
69        delayedRemovals.merge(num, 1, Integer::sum);
70      
71        // Update size and prune if necessary
72        if (num <= smallerHalf.peek()) {
73            --smallerHalfSize;
74            if (num == smallerHalf.peek()) {
75                prune(smallerHalf);
76            }
77        } else {
78            --largerHalfSize;
79            if (num == largerHalf.peek()) {
80                prune(largerHalf);
81            }
82        }
83      
84        // Maintain heap balance after removal
85        rebalance();
86    }
87
88    /**
89     * Removes all delayed elements from the top of the given heap
90     * @param heap the heap to prune
91     */
92    private void prune(PriorityQueue<Integer> heap) {
93        while (!heap.isEmpty() && delayedRemovals.containsKey(heap.peek())) {
94            // Decrement the delayed removal count
95            if (delayedRemovals.merge(heap.peek(), -1, Integer::sum) == 0) {
96                delayedRemovals.remove(heap.peek());
97            }
98            heap.poll();
99        }
100    }
101
102    /**
103     * Rebalances the two heaps to maintain the median property
104     * Ensures that smallerHalf has at most one more element than largerHalf
105     */
106    private void rebalance() {
107        // If smaller half has more than one extra element, move to larger half
108        if (smallerHalfSize > largerHalfSize + 1) {
109            largerHalf.offer(smallerHalf.poll());
110            --smallerHalfSize;
111            ++largerHalfSize;
112            prune(smallerHalf);
113        } 
114        // If larger half has more elements, move to smaller half
115        else if (smallerHalfSize < largerHalfSize) {
116            smallerHalf.offer(largerHalf.poll());
117            --largerHalfSize;
118            ++smallerHalfSize;
119            prune(largerHalf);
120        }
121    }
122}
123
124/**
125 * Solution class for the Sliding Window Median problem
126 */
127class Solution {
128    /**
129     * Calculates the median for each sliding window of size k in the given array
130     * @param nums the input array
131     * @param k the size of the sliding window
132     * @return array containing the median of each window
133     */
134    public double[] medianSlidingWindow(int[] nums, int k) {
135        // Initialize the MedianFinder with window size k
136        MedianFinder medianFinder = new MedianFinder(k);
137      
138        // Add the first k elements to initialize the window
139        for (int i = 0; i < k; ++i) {
140            medianFinder.addNum(nums[i]);
141        }
142      
143        // Prepare result array
144        int n = nums.length;
145        double[] result = new double[n - k + 1];
146      
147        // Calculate median for the first window
148        result[0] = medianFinder.findMedian();
149      
150        // Slide the window and calculate median for each position
151        for (int i = k; i < n; ++i) {
152            medianFinder.addNum(nums[i]);                    // Add new element
153            medianFinder.removeNum(nums[i - k]);             // Remove oldest element
154            result[i - k + 1] = medianFinder.findMedian();   // Store median
155        }
156      
157        return result;
158    }
159}
160
1class MedianFinder {
2public:
3    // Constructor: initialize with window size k
4    MedianFinder(int k) {
5        this->windowSize = k;
6    }
7
8    // Add a new number to the data structure
9    void addNum(int num) {
10        // Add to maxHeap if it's empty or num is less than or equal to max of maxHeap
11        if (maxHeap.empty() || num <= maxHeap.top()) {
12            maxHeap.push(num);
13            ++maxHeapSize;
14        } else {
15            // Otherwise add to minHeap
16            minHeap.push(num);
17            ++minHeapSize;
18        }
19        // Rebalance the two heaps to maintain size property
20        rebalance();
21    }
22
23    // Remove a number from the data structure (lazy deletion)
24    void removeNum(int num) {
25        // Mark the number for delayed removal
26        ++delayedRemoval[num];
27      
28        // Update size counter based on which heap the number belongs to
29        if (num <= maxHeap.top()) {
30            --maxHeapSize;
31            // If the number is at the top, prune immediately
32            if (num == maxHeap.top()) {
33                pruneHeap(maxHeap);
34            }
35        } else {
36            --minHeapSize;
37            // If the number is at the top, prune immediately
38            if (num == minHeap.top()) {
39                pruneHeap(minHeap);
40            }
41        }
42        // Rebalance the two heaps after removal
43        rebalance();
44    }
45
46    // Find and return the median of current numbers
47    double findMedian() {
48        // If window size is odd, return the top of maxHeap
49        // If even, return the average of tops of both heaps
50        return windowSize & 1 ? maxHeap.top() 
51                              : (static_cast<double>(maxHeap.top()) + minHeap.top()) / 2.0;
52    }
53
54private:
55    // Max heap to store the smaller half of numbers
56    priority_queue<int> maxHeap;
57  
58    // Min heap to store the larger half of numbers
59    priority_queue<int, vector<int>, greater<int>> minHeap;
60  
61    // Map to track numbers that need to be removed (lazy deletion)
62    unordered_map<int, int> delayedRemoval;
63  
64    // Actual size of maxHeap (excluding delayed removals)
65    int maxHeapSize = 0;
66  
67    // Actual size of minHeap (excluding delayed removals)
68    int minHeapSize = 0;
69  
70    // Window size
71    int windowSize;
72
73    // Remove delayed elements from the top of the heap
74    template <typename HeapType>
75    void pruneHeap(HeapType& heap) {
76        // Keep removing top elements that are marked for deletion
77        while (!heap.empty() && delayedRemoval[heap.top()]) {
78            // Decrement the delayed count
79            if (--delayedRemoval[heap.top()] == 0) {
80                // Remove from map if count reaches 0
81                delayedRemoval.erase(heap.top());
82            }
83            heap.pop();
84        }
85    }
86
87    // Rebalance the two heaps to maintain the invariant:
88    // maxHeap size should be equal to minHeap size (for even k)
89    // or maxHeap size should be minHeap size + 1 (for odd k)
90    void rebalance() {
91        // If maxHeap has more than one extra element, move top to minHeap
92        if (maxHeapSize > minHeapSize + 1) {
93            minHeap.push(maxHeap.top());
94            maxHeap.pop();
95            --maxHeapSize;
96            ++minHeapSize;
97            // Prune maxHeap after removing top
98            pruneHeap(maxHeap);
99        } 
100        // If minHeap has more elements, move top to maxHeap
101        else if (maxHeapSize < minHeapSize) {
102            maxHeap.push(minHeap.top());
103            minHeap.pop();
104            ++maxHeapSize;
105            --minHeapSize;
106            // Prune minHeap after removing top
107            pruneHeap(minHeap);
108        }
109    }
110};
111
112class Solution {
113public:
114    // Find medians of all k-sized sliding windows in the array
115    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
116        // Initialize MedianFinder with window size k
117        MedianFinder finder(k);
118      
119        // Add first k elements to build initial window
120        for (int i = 0; i < k; ++i) {
121            finder.addNum(nums[i]);
122        }
123      
124        // Store the median of first window
125        vector<double> result = {finder.findMedian()};
126      
127        // Slide the window through the rest of the array
128        for (int i = k; i < nums.size(); ++i) {
129            // Add new element entering the window
130            finder.addNum(nums[i]);
131            // Remove element leaving the window
132            finder.removeNum(nums[i - k]);
133            // Calculate and store median of current window
134            result.push_back(finder.findMedian());
135        }
136      
137        return result;
138    }
139};
140
1// Max heap implementation using array (stores smaller half of numbers)
2let maxHeap: number[] = [];
3// Min heap implementation using array (stores larger half of numbers)
4let minHeap: number[] = [];
5// Map to track numbers that need to be removed (lazy deletion)
6let delayedRemoval: Map<number, number> = new Map();
7// Actual size of maxHeap (excluding delayed removals)
8let maxHeapSize: number = 0;
9// Actual size of minHeap (excluding delayed removals)
10let minHeapSize: number = 0;
11// Window size
12let windowSize: number = 0;
13
14// Helper function to maintain max heap property after insertion
15function pushMaxHeap(num: number): void {
16    maxHeap.push(num);
17    let childIndex = maxHeap.length - 1;
18    while (childIndex > 0) {
19        const parentIndex = Math.floor((childIndex - 1) / 2);
20        if (maxHeap[parentIndex] >= maxHeap[childIndex]) break;
21        [maxHeap[parentIndex], maxHeap[childIndex]] = [maxHeap[childIndex], maxHeap[parentIndex]];
22        childIndex = parentIndex;
23    }
24}
25
26// Helper function to maintain min heap property after insertion
27function pushMinHeap(num: number): void {
28    minHeap.push(num);
29    let childIndex = minHeap.length - 1;
30    while (childIndex > 0) {
31        const parentIndex = Math.floor((childIndex - 1) / 2);
32        if (minHeap[parentIndex] <= minHeap[childIndex]) break;
33        [minHeap[parentIndex], minHeap[childIndex]] = [minHeap[childIndex], minHeap[parentIndex]];
34        childIndex = parentIndex;
35    }
36}
37
38// Helper function to maintain max heap property after removal
39function popMaxHeap(): void {
40    if (maxHeap.length === 0) return;
41    maxHeap[0] = maxHeap[maxHeap.length - 1];
42    maxHeap.pop();
43    if (maxHeap.length === 0) return;
44  
45    let parentIndex = 0;
46    while (true) {
47        let largest = parentIndex;
48        const leftChild = 2 * parentIndex + 1;
49        const rightChild = 2 * parentIndex + 2;
50      
51        if (leftChild < maxHeap.length && maxHeap[leftChild] > maxHeap[largest]) {
52            largest = leftChild;
53        }
54        if (rightChild < maxHeap.length && maxHeap[rightChild] > maxHeap[largest]) {
55            largest = rightChild;
56        }
57        if (largest === parentIndex) break;
58      
59        [maxHeap[parentIndex], maxHeap[largest]] = [maxHeap[largest], maxHeap[parentIndex]];
60        parentIndex = largest;
61    }
62}
63
64// Helper function to maintain min heap property after removal
65function popMinHeap(): void {
66    if (minHeap.length === 0) return;
67    minHeap[0] = minHeap[minHeap.length - 1];
68    minHeap.pop();
69    if (minHeap.length === 0) return;
70  
71    let parentIndex = 0;
72    while (true) {
73        let smallest = parentIndex;
74        const leftChild = 2 * parentIndex + 1;
75        const rightChild = 2 * parentIndex + 2;
76      
77        if (leftChild < minHeap.length && minHeap[leftChild] < minHeap[smallest]) {
78            smallest = leftChild;
79        }
80        if (rightChild < minHeap.length && minHeap[rightChild] < minHeap[smallest]) {
81            smallest = rightChild;
82        }
83        if (smallest === parentIndex) break;
84      
85        [minHeap[parentIndex], minHeap[smallest]] = [minHeap[smallest], minHeap[parentIndex]];
86        parentIndex = smallest;
87    }
88}
89
90// Initialize with window size k
91function initializeMedianFinder(k: number): void {
92    windowSize = k;
93    maxHeap = [];
94    minHeap = [];
95    delayedRemoval = new Map();
96    maxHeapSize = 0;
97    minHeapSize = 0;
98}
99
100// Remove delayed elements from the top of max heap
101function pruneMaxHeap(): void {
102    // Keep removing top elements that are marked for deletion
103    while (maxHeap.length > 0 && delayedRemoval.has(maxHeap[0]) && delayedRemoval.get(maxHeap[0])! > 0) {
104        const topValue = maxHeap[0];
105        const count = delayedRemoval.get(topValue)!;
106        // Decrement the delayed count
107        if (count === 1) {
108            // Remove from map if count reaches 0
109            delayedRemoval.delete(topValue);
110        } else {
111            delayedRemoval.set(topValue, count - 1);
112        }
113        popMaxHeap();
114    }
115}
116
117// Remove delayed elements from the top of min heap
118function pruneMinHeap(): void {
119    // Keep removing top elements that are marked for deletion
120    while (minHeap.length > 0 && delayedRemoval.has(minHeap[0]) && delayedRemoval.get(minHeap[0])! > 0) {
121        const topValue = minHeap[0];
122        const count = delayedRemoval.get(topValue)!;
123        // Decrement the delayed count
124        if (count === 1) {
125            // Remove from map if count reaches 0
126            delayedRemoval.delete(topValue);
127        } else {
128            delayedRemoval.set(topValue, count - 1);
129        }
130        popMinHeap();
131    }
132}
133
134// Rebalance the two heaps to maintain the invariant:
135// maxHeap size should be equal to minHeap size (for even k)
136// or maxHeap size should be minHeap size + 1 (for odd k)
137function rebalance(): void {
138    // If maxHeap has more than one extra element, move top to minHeap
139    if (maxHeapSize > minHeapSize + 1) {
140        pushMinHeap(maxHeap[0]);
141        popMaxHeap();
142        maxHeapSize--;
143        minHeapSize++;
144        // Prune maxHeap after removing top
145        pruneMaxHeap();
146    } 
147    // If minHeap has more elements, move top to maxHeap
148    else if (maxHeapSize < minHeapSize) {
149        pushMaxHeap(minHeap[0]);
150        popMinHeap();
151        maxHeapSize++;
152        minHeapSize--;
153        // Prune minHeap after removing top
154        pruneMinHeap();
155    }
156}
157
158// Add a new number to the data structure
159function addNum(num: number): void {
160    // Add to maxHeap if it's empty or num is less than or equal to max of maxHeap
161    if (maxHeap.length === 0 || num <= maxHeap[0]) {
162        pushMaxHeap(num);
163        maxHeapSize++;
164    } else {
165        // Otherwise add to minHeap
166        pushMinHeap(num);
167        minHeapSize++;
168    }
169    // Rebalance the two heaps to maintain size property
170    rebalance();
171}
172
173// Remove a number from the data structure (lazy deletion)
174function removeNum(num: number): void {
175    // Mark the number for delayed removal
176    delayedRemoval.set(num, (delayedRemoval.get(num) || 0) + 1);
177  
178    // Update size counter based on which heap the number belongs to
179    if (num <= maxHeap[0]) {
180        maxHeapSize--;
181        // If the number is at the top, prune immediately
182        if (num === maxHeap[0]) {
183            pruneMaxHeap();
184        }
185    } else {
186        minHeapSize--;
187        // If the number is at the top, prune immediately
188        if (num === minHeap[0]) {
189            pruneMinHeap();
190        }
191    }
192    // Rebalance the two heaps after removal
193    rebalance();
194}
195
196// Find and return the median of current numbers
197function findMedian(): number {
198    // If window size is odd, return the top of maxHeap
199    // If even, return the average of tops of both heaps
200    return (windowSize & 1) ? maxHeap[0] : (maxHeap[0] + minHeap[0]) / 2.0;
201}
202
203// Find medians of all k-sized sliding windows in the array
204function medianSlidingWindow(nums: number[], k: number): number[] {
205    // Initialize MedianFinder with window size k
206    initializeMedianFinder(k);
207  
208    // Add first k elements to build initial window
209    for (let i = 0; i < k; i++) {
210        addNum(nums[i]);
211    }
212  
213    // Store the median of first window
214    const result: number[] = [findMedian()];
215  
216    // Slide the window through the rest of the array
217    for (let i = k; i < nums.length; i++) {
218        // Add new element entering the window
219        addNum(nums[i]);
220        // Remove element leaving the window
221        removeNum(nums[i - k]);
222        // Calculate and store median of current window
223        result.push(findMedian());
224    }
225  
226    return result;
227}
228

Time and Space Complexity

Time Complexity: O(n × log k)

The algorithm uses a sliding window approach with two heaps to maintain the median. For each of the n elements in the array:

  • add_num(): Inserting into a heap takes O(log k) time since the heap size is bounded by k
  • remove_num(): Uses lazy deletion - marking for deletion is O(1), but pruning may trigger multiple heap pops in worst case
  • prune(): In amortized analysis, each element is popped at most once, contributing O(log k) per element
  • rebalance(): Involves at most one heap pop and push operation, which is O(log k)
  • find_median(): Returns the top element(s) of heaps in O(1) time

Since we process n elements and each operation involves heap operations on heaps of size at most k, the total time complexity is O(n × log k).

Note: The reference answer states O(n × log n), which would be the case if the heap size could grow to n. However, since the sliding window size is k, the heap operations are actually O(log k), making the overall complexity O(n × log k).

Space Complexity: O(n)

  • The two heaps (small and large) together store at most k elements at any time: O(k)
  • The delayed dictionary in worst case could store all unique elements from the array if they all get delayed for deletion: O(n)
  • The result array ans stores n - k + 1 median values: O(n)

The dominant space usage is O(n) from the delayed dictionary and result array.

Common Pitfalls

1. Incorrect Heap Selection During Removal

Pitfall: When removing an element, determining which heap it belongs to by comparing with the current top of the small heap can be incorrect if the top itself is a delayed element that hasn't been pruned yet.

Example Scenario:

  • Small heap: [-5, -3] (actual values: [5, 3])
  • Large heap: [6, 8]
  • Delayed: {5: 1} (5 is marked for deletion but still at top)
  • Trying to remove 4: Since 4 <= -(-5) = 5, we'd incorrectly think 4 is in the small heap

Solution: Always prune the small heap before comparing:

def remove_num(self, num: int) -> None:
    self.delayed[num] += 1
  
    # Prune first to ensure valid comparison
    self.prune(self.small)
  
    if self.small and num <= -self.small[0]:
        self.small_size -= 1
        if num == -self.small[0]:
            self.prune(self.small)
    else:
        self.large_size -= 1
        if self.large and num == self.large[0]:
            self.prune(self.large)
  
    self.rebalance()

2. Integer Overflow in Median Calculation

Pitfall: For very large integers, calculating (-small[0] + large[0]) / 2 might cause overflow in languages with fixed integer sizes.

Solution: Use floating-point division from the start:

def find_median(self) -> float:
    if self.k & 1:
        return float(-self.small[0])
    else:
        return (float(-self.small[0]) + float(self.large[0])) / 2.0

3. Not Handling Edge Cases in Rebalancing

Pitfall: The rebalance function might try to pop from an empty heap after pruning, especially when all elements in a heap are marked for deletion.

Solution: Add safety checks:

def rebalance(self) -> None:
    if self.small_size > self.large_size + 1:
        # Ensure small heap has valid elements before popping
        self.prune(self.small)
        if self.small:
            heappush(self.large, -heappop(self.small))
            self.small_size -= 1
            self.large_size += 1
            self.prune(self.small)
    elif self.small_size < self.large_size:
        # Ensure large heap has valid elements before popping
        self.prune(self.large)
        if self.large:
            heappush(self.small, -heappop(self.large))
            self.large_size -= 1
            self.small_size += 1
            self.prune(self.large)

4. Inefficient Pruning Strategy

Pitfall: Only pruning when necessary (at heap tops) can lead to heap bloat with many delayed elements accumulating in the middle of heaps, causing increased memory usage and slower heap operations.

Solution: Implement periodic full cleanup or track a threshold:

def add_num(self, num: int) -> None:
    # Existing add logic...
  
    # Periodic cleanup when delayed elements exceed threshold
    if len(self.delayed) > self.k // 2:
        self.full_cleanup()
  
    self.rebalance()

def full_cleanup(self) -> None:
    # Rebuild heaps without delayed elements
    valid_nums = []
    while self.small:
        val = -heappop(self.small)
        if val not in self.delayed:
            valid_nums.append(val)
    while self.large:
        val = heappop(self.large)
        if val not in self.delayed:
            valid_nums.append(val)
  
    # Clear delayed dictionary and rebuild
    self.delayed.clear()
    self.small = []
    self.large = []
    self.small_size = 0
    self.large_size = 0
  
    for num in valid_nums:
        self.add_num(num)

5. Duplicate Values Handling

Pitfall: When the same value appears multiple times and some instances are delayed while others are active, the lazy deletion counting can become inconsistent.

Solution: The current implementation handles this correctly by using a count-based approach in the delayed dictionary, but ensure the count never goes negative:

def prune(self, heap: List[int]) -> None:
    sign = -1 if heap is self.small else 1
  
    while heap and sign * heap[0] in self.delayed:
        # Ensure we don't over-decrement
        if self.delayed[sign * heap[0]] > 0:
            self.delayed[sign * heap[0]] -= 1
            if self.delayed[sign * heap[0]] == 0:
                self.delayed.pop(sign * heap[0])
        heappop(heap)
Discover Your Strengths and Weaknesses: Take Our 3-Minute Quiz to Tailor Your Study Plan:

Which of the two traversal algorithms (BFS and DFS) can be used to find whether two nodes are connected?

Recommended Readings

Sliding Window Maximum

https assets algo monster cover_photos stack svg Sliding Window Maximum Monotonic Stack We have an array and a sliding window defined by a start index and an end index The sliding window moves from left of the array to right There are always k elements in the window The window

Heap Fundamentals

https assets algo monster cover_photos heap svg Priority Queue and Heap What is the relationship between priority queue and heap Priority Queue is an Abstract Data Type and Heap is the concrete data structure we use to implement a priority queue Priority Queue A priority queue is a data structure

Coding Interview Patterns: The Smart, Data-Driven Way to Prepare for Tech Interviews

Coding Interview Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More