Sliding Window Median

Intuition

Solution

class Solution:
    def pop(self, heap, valid):
        """
        :param heap: each item in heap is a pair (value, index) where index is
                     the index in list nums
        :param valid: minimal index which is valid
        """
        index = valid - 1
        while index < valid:
            value, index = heapq.heappop(heap)
        return value, index

    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        """
        Two heap
        Almost the same problem as 295. Find Median from data stream.
        """
        medians = []
        left = []
        left_len = 0
        right = []
        right_len = 0
        in_left = [False] * len(nums)

        for index, num in enumerate(nums):
            valid = index - k + 1
            # Always keep right >= left
            if left_len < right_len:
                heapq.heappush(right, (num, index))
                v, i = self.pop(right, valid)
                heapq.heappush(
                    left,
                    (-v, i)
                )
                left_len += 1
                in_left[index] = False
                in_left[i] = True
            else:
                heapq.heappush(left, (-num, index))
                v, i = self.pop(left, valid)
                heapq.heappush(
                    right,
                    (-v, i)
                )
                right_len += 1
                in_left[index] = True
                in_left[i] = False

            if left_len + right_len == k:
                while left and left[0][1] < valid:
                    heapq.heappop(left)
                while right and right[0][1] < valid:
                    heapq.heappop(right)

                if left_len == right_len:
                    medians.append((-left[0][0] + right[0][0]) / 2)
                else:
                    medians.append(right[0][0])

                # Remove first in the heap
                # Need to use lazy delete here because python does not support
                # remove
                if in_left[index - k + 1]:
                    left_len -= 1
                else:
                    right_len -= 1
        return medians