LeetCode Solutions Blog

Sliding Window Median

Number: 480

Difficulty: Hard

Paid? No

Problem Description

Given an integer array nums and an integer k representing the size of a sliding window, compute the median of each window as it moves from the beginning to the end of the array. The median is defined as the middle value in a sorted list of numbers; if there is an even number of elements, it is the mean of the two middle values.

Key Insights

Use the sliding window technique to sequentially add and remove elements.
Maintain two heaps:
- A max heap (implemented via negatives) for the lower half of the data.
- A min heap for the higher half of the data.
The two heaps allow efficient median retrieval in O(1) time.
Lazy removal and balancing of heaps are crucial to ensure the correctness when elements exit the window.

Space and Time Complexity

Time Complexity: O(n log k) because each insertion and removal in the heaps takes O(log k). Space Complexity: O(k) to maintain the two heaps.

Solution

We solve the problem using a DualHeap approach. We maintain two heaps to divide the sliding window into lower and upper halves. When a new element is added, it is placed into one of the heaps based on its value relative to the current median. When an element falls out of the window, we mark it for lazy removal in the appropriate heap. We then balance the heaps to ensure the sizes differ by at most one, allowing quick computation of the median.

Code Solutions

Below are sample implementations in multiple languages.

import heapq

class DualHeap:
    def __init__(self, k):
        self.small = []  # max heap (store negatives)
        self.large = []  # min heap
        self.k = k
        self.delayed = {}  # counts for lazy removal
        self.smallSize = 0  # effective size of small heap
        self.largeSize = 0  # effective size of large heap

    def prune(self, heap):
        # Remove elements from the top that are marked for lazy deletion.
        while heap:
            num = heap[0] if heap is self.large else -heap[0]
            if num in self.delayed and self.delayed[num] > 0:
                heapq.heappop(heap)
                self.delayed[num] -= 1
                if self.delayed[num] == 0:
                    del self.delayed[num]
            else:
                break

    def balance(self):
        # Ensure the two heaps' sizes differ by at most one.
        if self.smallSize > self.largeSize + 1:
            val = -heapq.heappop(self.small)
            self.smallSize -= 1
            heapq.heappush(self.large, val)
            self.largeSize += 1
            self.prune(self.small)
        elif self.smallSize < self.largeSize:
            val = heapq.heappop(self.large)
            self.largeSize -= 1
            heapq.heappush(self.small, -val)
            self.smallSize += 1
            self.prune(self.large)

    def add(self, num):
        if not self.small or num <= -self.small[0]:
            heapq.heappush(self.small, -num)
            self.smallSize += 1
        else:
            heapq.heappush(self.large, num)
            self.largeSize += 1
        self.balance()

    def remove(self, num):
        # Mark the element for lazy removal.
        self.delayed[num] = self.delayed.get(num, 0) + 1
        if self.small and num <= -self.small[0]:
            self.smallSize -= 1
            if num == -self.small[0]:
                self.prune(self.small)
        else:
            self.largeSize -= 1
            if self.large and num == self.large[0]:
                self.prune(self.large)
        self.balance()

    def median(self):
        if self.k % 2:
            return float(-self.small[0])
        else:
            return (-self.small[0] + self.large[0]) / 2.0

def slidingWindowMedian(nums, k):
    result = []
    dh = DualHeap(k)
    for i, num in enumerate(nums):
        dh.add(num)
        if i >= k - 1:
            result.append(dh.median())
            dh.remove(nums[i - k + 1])
    return result

# Example usage:
nums = [1, 3, -1, -3, 5, 3, 6, 7]
k = 3
print(slidingWindowMedian(nums, k))