Evgeny's Notes

Recent posts

480. Sliding Window Median

Dec 19, 20243 min read

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle values.

  • For examples, if arr = [2,3,4], the median is 3.
  • For examples, if arr = [1,2,3,4], the median is (2 + 3) / 2 = 2.5.

You are given an integer array nums and an integer k. There is a sliding window of size k which is moving from the very left of the array to the very right. You can only see the k numbers in the window. Each time the sliding window moves right by one position.

Return the median array for each window in the original array. Answers within 10-5 of the actual value will be accepted.

import heapq
 
class Solution:
    def medianSlidingWindow(self, nums, k):
        # Two heaps: small (max-heap), large (min-heap)
        small, large = [], []
        
        # Initialize the first window
        for i, x in enumerate(nums[:k]):
            heapq.heappush(small, (-x, i))
        
        # Balance the heaps: ensure large has at least (k//2) elements
        for _ in range(k - (k // 2)):
            self.move(small, large)
        
        # Store the medians
        result = [self.get_median(small, large, k)]
        
        # Process the sliding window
        for i, x in enumerate(nums[k:], start=k):
            # Add the new element to the appropriate heap
            if x >= large[0][0]:
                heapq.heappush(large, (x, i))
                if nums[i - k] <= large[0][0]:
                    self.move(large, small)
            else:
                heapq.heappush(small, (-x, i))
                if nums[i - k] >= large[0][0]:
                    self.move(small, large)
            
            # Remove elements outside the current window
            while small and small[0][1] <= i - k:
                heapq.heappop(small)
            while large and large[0][1] <= i - k:
                heapq.heappop(large)
            
            # Append the current median
            result.append(self.get_median(small, large, k))
        
        return result
 
    def move(self, source_heap, target_heap):
        """Move the top element from source_heap to target_heap."""
        value, index = heapq.heappop(source_heap)
        heapq.heappush(target_heap, (-value, index))
    
    def get_median(self, small, large, k):
        """Calculate the median based on the heaps."""
        if k % 2 == 1:  # Odd window size
            return float(large[0][0])
        else:  # Even window size
            return (large[0][0] - small[0][0]) / 2.0

Code Explanation

This code calculates the median of a sliding window of size k over the input list nums.

1. Two Heaps (small and large)

  • The algorithm uses two heaps:
    • small: A max-heap (simulated using a min-heap with negative values).
    • large: A min-heap.
  • Together, they partition the sliding window such that:
    • small contains the smaller half of the numbers (as negatives for max-heap behavior).
    • large contains the larger half of the numbers.

2. Initial Window Setup

  • The first k elements of nums are pushed into the small heap.
  • To balance the heaps, elements are moved from small to large until large contains at least k/2 elements (for even k) or one extra element for odd k.

3. Sliding Window Processing

  • As the window slides through nums, the algorithm:
    1. Adds the next number to one of the heaps, depending on its value relative to the current median.
    2. Moves an element between heaps if necessary to maintain balance.
    3. Removes outdated elements (outside the current window) from the heaps.
    4. Computes the median after each step and appends it to the result list.

4. Helper Functions

  • move(h1, h2): Moves the top element from heap h1 to h2. Adjusts signs for maintaining the max-heap in small.
  • get_med(h1, h2, k): Computes the median based on the top elements of the heaps.

Graph View

Backlinks