Problem statement

https://leetcode.com/problems/sliding-window-median/

Solution

This problem is similar to problem 0295. Find Median from Data Stream. We need to keep two heaps: one for small half and another for large half. Also we need to perform so-called lazy deletions from our heap: we check each iteration the top of our heaps and if we have old elements, we delete them.

Complexity

Time complexity is $O(n\log n)$, because potentially we can have upto $n$ elements in heaps due to lazy updates. Space complexity is $O(n)$.

Code

class Solution:
    def medianSlidingWindow(self, nums, k):
        lo, hi, ans = [], [], []
        nums = [0] + nums
        
        def send(h1, h2):  #sends element from heap h1 to heap h2 .
            heappush(h2, (-h1[0][0], h1[0][1]))
            heappop(h1)

        for i, n in enumerate(nums[:k]): heappush(lo, (-n, i))
        for i in range((k+1)//2): send(lo, hi)
        
        for i, num in enumerate(nums[k:]):
            if num >= hi[0][0]:
                heappush(hi, (num, i+k))
                if nums[i] <= hi[0][0]: send(hi, lo)
            else:
                heappush(lo, (-num, i+k))
                if nums[i] >= hi[0][0]: send(lo, hi)    

            while(lo and lo[0][1] <= i): heappop(lo)  # lazy-deletion
            while(hi and hi[0][1] <= i): heappop(hi)  # lazy-deletion
                
            ans.append(hi[0][0] if k%2 else (hi[0][0] - lo[0][0])/2)
            
        return ans