leetcode-480 Sliding Window Median

# The median is the middle value in an ordered list. Return the median array
# for each window of size k.

from heapq import *

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        # Two heaps: max_heap for smaller half, min_heap for larger half
        max_heap, min_heap = [], []
        result = []

        for i, num in enumerate(nums):
            # Add to appropriate heap
            if not max_heap or num <= -max_heap[0]:
                heappush(max_heap, -num)
            else:
                heappush(min_heap, num)

            # Balance heaps
            if len(max_heap) > len(min_heap) + 1:
                heappush(min_heap, -heappop(max_heap))
            elif len(min_heap) > len(max_heap):
                heappush(max_heap, -heappop(min_heap))

            # Get median when window is full
            if i >= k - 1:
                if k % 2:
                    result.append(float(-max_heap[0]))
                else:
                    result.append((-max_heap[0] + min_heap[0]) / 2.0)

                # Remove leftmost element (simplified - full solution needs lazy deletion)
                to_remove = nums[i - k + 1]

        return result

# Time: O(n log k), Space: O(k)
# AirBnB: Advanced heap problem with sliding window