Hard

Sliding Window Median

LAB

Return the median for every contiguous window of size k.

EXAMPLES

Example 1

Input
{
  "nums": [
    1,
    3,
    -1,
    -3,
    5,
    3,
    6,
    7
  ],
  "k": 3
}

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

FUNCTION SHAPE

nums: intArrayk: int→doubleArray

SOLUTION NOTE

Uses 2 SortedLists instead of 2 heaps because heaps do not support O(log n) deletion. For a MedianFinder that needs removals, we need SortedLists.

Reveal reference solution +

pythonREFERENCE

from sortedcontainers import SortedList

class MedianFinder:
    def __init__(self):
        self.left = SortedList()   # smaller half (or one extra)
        self.right = SortedList()  # larger half

    def _rebalance(self):
        if len(self.left) < len(self.right):
            self.left.add(self.right.pop(0))
        elif len(self.left) > len(self.right) + 1:
            self.right.add(self.left.pop(-1))

    def addNum(self, num: int) -> None:
        self.left.add(num)
        self.right.add(self.left.pop(-1))
        self._rebalance()

    def removeNum(self, num: int) -> None:
        try:
            self.left.remove(num)
        except ValueError:
            self.right.remove(num)
        self._rebalance()

    def findMedian(self) -> float:
        if len(self.left) > len(self.right):
            return float(self.left[-1])
        return (self.left[-1] + self.right[0]) / 2

def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
    mf = MedianFinder()
    n = len(nums)
    res = []
    for i in range(n):
        mf.addNum(nums[i])
        if i >= k - 1:
            res.append(mf.findMedian())
            mf.removeNum(nums[i - (k - 1)])
    return res

TimeO(n log k)

SpaceO(k)

Open on LeetCode ↗