480. Sliding Window Median - Leetcode Solution

Problem Description

Given an array of integers nums and an integer k, your task is to compute the median for every contiguous subarray (or "window") of size k as the window slides from the start to the end of the array. The output should be a list of the medians for each window, in the order they appear.

The median is defined as the middle value in an ordered list. If the window size k is odd, it is the middle element. If k is even, it is the average of the two middle elements.
Each window contains k consecutive elements from nums.
You must calculate the median for every possible window as it slides from left to right.
Constraints:
- 1 <= k <= nums.length <= 10⁵
- There is only one valid median for each window.
- Elements may repeat, but each window is distinct by its position.

Thought Process

The straightforward way to solve this problem is to iterate through each window, sort its elements, and pick the median. However, sorting each window from scratch would be too slow for large arrays because sorting takes O(k log k) time and there are O(n) windows.

To optimize, we need a data structure that allows us to efficiently insert and remove elements (as the window slides), and also find the median quickly. This leads us to consider using two heaps (a max-heap and a min-heap) or a balanced binary search tree (or multiset in C++). Heaps are especially well-suited because they allow us to keep track of the lower and upper halves of the window.

The key insight is to maintain the current window in a way that supports:

Efficient median calculation (O(1) or O(log k))
Efficient insertion and deletion as the window slides (O(log k))

Solution Approach

We solve the Sliding Window Median problem using a "two heaps" approach:

Use two heaps:
- A max-heap (small) for the smaller half of the numbers (Python's heapq is a min-heap, so we insert negatives to simulate a max-heap).
- A min-heap (large) for the larger half of the numbers.
Balance: Maintain the heaps so that their sizes differ by at most one. The max-heap can have one more element when k is odd.
Median Extraction:
- If k is odd, the median is the top of the max-heap.
- If k is even, the median is the average of the tops of both heaps.
Sliding the Window:
- When moving the window, remove the outgoing element and add the incoming one.
- Since heaps do not support efficient removal of arbitrary elements, we use a "delayed deletion" technique: mark elements to be deleted, and lazily remove them when they reach the top of the heap.
Repeat: For each window, record the median after balancing the heaps.

This approach ensures that each insertion, removal, and median calculation is O(log k), making it efficient for large inputs.

Example Walkthrough

Let's consider nums = [1, 3, -1, -3, 5, 3, 6, 7] and k = 3.

First window: [1, 3, -1] → Sorted: [-1, 1, 3], median is 1.
Slide right: [3, -1, -3] → Sorted: [-3, -1, 3], median is -1.
Slide right: [-1, -3, 5] → Sorted: [-3, -1, 5], median is -1.
Slide right: [-3, 5, 3] → Sorted: [-3, 3, 5], median is 3.
Slide right: [5, 3, 6] → Sorted: [3, 5, 6], median is 5.
Slide right: [3, 6, 7] → Sorted: [3, 6, 7], median is 6.

Output: [1, -1, -1, 3, 5, 6]

With the heap approach, we insert and remove elements while maintaining balance and efficiently extracting the median at each step.

Time and Space Complexity

Brute-force method:
- For each of the O(n) windows, sorting takes O(k log k).
- Total time: O(nk log k).
Optimized heap-based solution:
- Each insertion and removal from a heap is O(log k).
- Finding the median is O(1).
- Total time: O(n log k).
- Space: O(k) for the heaps and O(n) for the output.

The heap-based approach is much more efficient for large arrays and window sizes.

Summary

The Sliding Window Median problem asks us to efficiently compute the median of every window of size k in an array. The naive solution is too slow, so we use two heaps to keep track of the current window's lower and upper halves. By carefully balancing the heaps and using lazy deletion, we can insert, remove, and find medians in logarithmic time per operation. This method is both elegant and practical, making it suitable for large datasets and real-time scenarios.

Code Implementation

import heapq
from collections import defaultdict

class DualHeap:
    def __init__(self, k):
        self.small = []  # Max heap (invert values)
        self.large = []  # Min heap
        self.delayed = defaultdict(int)
        self.k = k
        self.small_size = 0
        self.large_size = 0

    def prune(self, heap):
        while heap:
            num = -heap[0] if heap == self.small else heap[0]
            if self.delayed[num]:
                heapq.heappop(heap)
                self.delayed[num] -= 1
            else:
                break

    def balance(self):
        # Balance sizes
        if self.small_size > self.large_size + 1:
            heapq.heappush(self.large, -heapq.heappop(self.small))
            self.small_size -= 1
            self.large_size += 1
            self.prune(self.small)
        elif self.small_size < self.large_size:
            heapq.heappush(self.small, -heapq.heappop(self.large))
            self.small_size += 1
            self.large_size -= 1
            self.prune(self.large)

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

    def erase(self, num):
        self.delayed[num] += 1
        if 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.balance()

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

def medianSlidingWindow(nums, k):
    dh = DualHeap(k)
    result = []
    for i in range(k):
        dh.insert(nums[i])
    result.append(dh.get_median())
    for i in range(k, len(nums)):
        dh.insert(nums[i])
        dh.erase(nums[i - k])
        result.append(dh.get_median())
    return result

#include <vector>
#include <queue>
#include <unordered_map>

using namespace std;

class DualHeap {
public:
    priority_queue<int> small; // max heap
    priority_queue<int, vector<int>, greater<int> > large; // min heap
    unordered_map<int, int> delayed;
    int k, smallSize = 0, largeSize = 0;

    DualHeap(int k): k(k) {}

    void prune(priority_queue<int> &heap) {
        while (!heap.empty()) {
            int num = heap.top();
            if (delayed[num]) {
                heap.pop();
                delayed[num]--;
            } else break;
        }
    }
    void prune(priority_queue<int, vector<int>, greater<int> > &heap) {
        while (!heap.empty()) {
            int num = heap.top();
            if (delayed[num]) {
                heap.pop();
                delayed[num]--;
            } else break;
        }
    }

    void balance() {
        if (smallSize > largeSize + 1) {
            large.push(small.top());
            small.pop();
            smallSize--;
            largeSize++;
            prune(small);
        } else if (smallSize < largeSize) {
            small.push(large.top());
            large.pop();
            smallSize++;
            largeSize--;
            prune(large);
        }
    }

    void insert(int num) {
        if (small.empty() || num <= small.top()) {
            small.push(num);
            smallSize++;
        } else {
            large.push(num);
            largeSize++;
        }
        balance();
    }

    void erase(int num) {
        delayed[num]++;
        if (!small.empty() && num <= small.top()) {
            smallSize--;
            if (num == small.top()) prune(small);
        } else {
            largeSize--;
            if (!large.empty() && num == large.top()) prune(large);
        }
        balance();
    }

    double getMedian() {
        if (k % 2) return small.top();
        else return ((double)small.top() + large.top()) / 2.0;
    }
};

vector<double> medianSlidingWindow(vector<int>& nums, int k) {
    DualHeap dh(k);
    vector<double> result;
    for (int i = 0; i < k; ++i) dh.insert(nums[i]);
    result.push_back(dh.getMedian());
    for (int i = k; i < nums.size(); ++i) {
        dh.insert(nums[i]);
        dh.erase(nums[i - k]);
        result.push_back(dh.getMedian());
    }
    return result;
}

import java.util.*;

class DualHeap {
    // Max heap for the smaller half, min heap for the larger half
    private PriorityQueue<Integer> small;
    private PriorityQueue<Integer> large;
    private Map<Integer, Integer> delayed;
    private int k, smallSize, largeSize;

    public DualHeap(int k) {
        this.k = k;
        this.small = new PriorityQueue<>(Collections.reverseOrder());
        this.large = new PriorityQueue<>();
        this.delayed = new HashMap<>();
        this.smallSize = 0;
        this.largeSize = 0;
    }

    private void prune(PriorityQueue<Integer> heap) {
        while (!heap.isEmpty()) {
            int num = heap.peek();
            if (delayed.getOrDefault(num, 0) > 0) {
                delayed.put(num, delayed.get(num) - 1);
                heap.poll();
            } else break;
        }
    }

    private void balance() {
        if (smallSize > largeSize + 1) {
            large.offer(small.poll());
            smallSize--;
            largeSize++;
            prune(small);
        } else if (smallSize < largeSize) {
            small.offer(large.poll());
            smallSize++;
            largeSize--;
            prune(large);
        }
    }

    public void insert(int num) {
        if (small.isEmpty() || num <= small.peek()) {
            small.offer(num);
            smallSize++;
        } else {
            large.offer(num);
            largeSize++;
        }
        balance();
    }

    public void erase(int num) {
        delayed.put(num, delayed.getOrDefault(num, 0) + 1);
        if (!small.isEmpty() && num <= small.peek()) {
            smallSize--;
            if (num == small.peek()) prune(small);
        } else {
            largeSize--;
            if (!large.isEmpty() && num == large.peek()) prune(large);
        }
        balance();
    }

    public double getMedian() {
        if (k % 2 == 1) {
            return (double) small.peek();
        } else {
            return ((double) small.peek() + large.peek()) / 2.0;
        }
    }
}

class Solution {
    public double[] medianSlidingWindow(int[] nums, int k) {
        DualHeap dh = new DualHeap(k);
        double[] result = new double[nums.length - k + 1];
        for (int i = 0; i < k; ++i) dh.insert(nums[i]);
        result[0] = dh.getMedian();
        for (int i = k; i < nums.length; ++i) {
            dh.insert(nums[i]);
            dh.erase(nums[i - k]);
            result[i - k + 1] = dh.getMedian();
        }
        return result;
    }
}

class Heap {
    constructor(compare) {
        this.data = [];
        this.compare = compare;
    }
    push(val) {
        this.data.push(val);
        this._siftUp();
    }
    pop() {
        const top = this.data[0];
        const last = this.data.pop();
        if (this.data.length) {
            this.data[0] = last;
            this._siftDown();
        }
        return top;
    }
    peek() {
        return this.data[0];
    }
    size() {
        return this.data.length;
    }
    _siftUp() {
        let i = this.data.length - 1;
        while (i > 0) {
            let p = (i - 1) >> 1;
            if (this.compare(this.data[i], this.data[p])) {
                [this.data[i], this.data[p]] = [this.data[p], this.data[i]];
                i = p;
            } else break;
        }
    }
    _siftDown() {
        let i = 0, n = this.data.length;
        while (true) {
            let l = i * 2 + 1, r = i * 2 + 2, next = i;
            if (l < n && this.compare(this.data[l], this.data[next])) next = l;
            if (r < n && this.compare(this.data[r], this.data[next])) next = r;
            if (next !== i) {
                [this.data[i], this.data[next]] = [this.data[next], this.data[i]];
                i = next;
            } else break;
        }
    }
}

function medianSlidingWindow(nums, k) {
    let small = new Heap((a, b) => a > b); // max heap
    let large = new Heap((a, b) => a < b); // min heap
    let delayed = new Map();

    function prune(heap) {
        while (heap.size()) {
            let num = heap.peek();
            if (delayed.get(num)) {
                heap.pop();
                delayed.set(num, delayed.get(num) - 1);
                if (delayed.get(num) === 0) delayed.delete(num);
            } else break;
        }
    }

    function balance() {
        if (small.size() > large.size() + 1) {
            large.push(small.pop());
            prune(small);
        } else if (small.size() < large.size()) {
            small.push(large.pop());
            prune(large);
        }
    }

    function insert(num) {
        if (!small.size() || num <= small.peek()) small.push(num);
        else large.push(num);
        balance();
    }

    function erase(num) {
        delayed.set(num, (delayed.get(num) || 0) + 1);
        if (num <= small.peek()) {
            prune(small);
        } else {
            prune(large);
        }
        balance();
    }

    function getMedian() {
        if (k % 2) return small.peek();
        else return (small.peek() + large.peek()) / 2;
    }

    let result = [];
    for (let i = 0; i < k; ++i) insert(nums[i]);
    result.push(getMedian());
    for (let i = k; i < nums.length; ++i) {
        insert(nums[i]);
        erase(nums[i - k]);
        result.push(getMedian());
    }
    return result;
}

AlgoMap.io

Want Help Cracking FAANG?