LeetCode 480: Sliding Window Median

Like the data stream median problem, we use two heaps.

However, as numbers fall off the sliding window, we should also remove them from the heaps. Since, heaps do not support efficient search and delete, we delete lazily. When an index fall off the sliding window, we mark that index as stale. In the operations that look at tops, like push() or median(), we first get rid of stale tops.

If the numbers are in random order, the heap size should be about the size of the sliding window. So, time: $\mathcal{O}(n \cdot \lg{k})$ , space: $\mathcal{O}(k)$ .

However, if the numbers are sorted in increasing order, in the left_maxq, the stale indices will never surface to the top, so the size of that heap will grow like $\mathcal{O}(n)$ making the heap operations slower, taking time $\mathcal{O}(n)$ . In such worst-case, time: $\mathcal{O}(n \cdot \lg{n})$ , space: $\mathcal{O}(n)$ .

from heapq import heappush as push, heappop as pop

class MedianFinder:
    def __init__(self):
        self.left_size, self.right_size = 0, 0
        self.left_indices, self.right_indices = set(), set()
        self.left_stale_indices, self.right_stale_indices = set(), set()
        self.left_maxq, self.right_minq = [], []

    def __delete_stale_top(self):
        while self.left_maxq and (self.left_maxq[0][1] in self.left_stale_indices):
            self.left_stale_indices.remove(self.left_maxq[0][1])
            pop(self.left_maxq)
        
        while self.right_minq and (self.right_minq[0][1] in self.right_stale_indices):
            self.right_stale_indices.remove(self.right_minq[0][1])
            pop(self.right_minq)
    
    def push(self, index, num):
        self.__delete_stale_top()

        if self.left_size == 0 or num <= -self.left_maxq[0][0]:
            self.left_size += 1
            push(self.left_maxq, (-num, index))
            self.left_indices.add(index)
        else:
            self.right_size += 1
            push(self.right_minq, (num, index))
            self.right_indices.add(index)

        if self.left_size < self.right_size:
            x, i = pop(self.right_minq)
            self.right_size -= 1
            self.right_indices.remove(i)
            
            push(self.left_maxq, (-x, i))
            self.left_size += 1
            self.left_indices.add(i)
        elif self.left_size - self.right_size > 1:
            neg_x, i = pop(self.left_maxq)
            self.left_size -= 1
            self.left_indices.remove(i)
            
            push(self.right_minq, (-neg_x, i))
            self.right_size += 1
            self.right_indices.add(i)

    def mark_stale(self, index):
        if index in self.left_indices:
            self.left_size -= 1
            self.left_indices.remove(index)
            self.left_stale_indices.add(index)
        else:
            self.right_size -= 1
            self.right_indices.remove(index)
            self.right_stale_indices.add(index)

    def median(self):
        self.__delete_stale_top()

        if self.left_size > self.right_size:
            return -self.left_maxq[0][0]
        else:
            return 0.5 * ( -self.left_maxq[0][0] + self.right_minq[0][0] )

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        medians = []
        median_finder = MedianFinder()
        for i, x in enumerate(nums):
            if i >= k:
                median_finder.mark_stale(i-k)
            median_finder.push(i, x)
            if i >= k-1:
                medians.append(median_finder.median())

        return medians

Updatable Heap

It is possible to support efficient change-key operation on heap. In that case, for the index that falls off the sliding window, we could change its key or the associated number so it surfaces to the top of the heap and then we could pop it. In other words, the remove will take $\mathcal{O}(\lg{k})$ time and the heap sizes will be bounded at $\mathcal{O}(k)$ .

unreasonably effective