480. Sliding Window Median (Array, Hash Table, Sliding Window, Heap (Priority Queue))

링크 - https://leetcode.com/problems/sliding-window-median/description/?envType=company&envId=facebook&favoriteSlug=facebook-thirty-days

난이도 - Hard

Intuition

이 문제는 min heap(large)과 max heap(small)두개를 사용해서 풀어야하는 문제이다. 어려운 점은 계속 이전에 맨 왼쪽에 있는 원소를 제거해야 한다는 점이다. 따라서 이를 lazy removal방식으로 해결했다.

1. 처음에 k개의 원소를 heap들에 넣어준다. small과 large의 개수가 같을때는, small에 먼저 넣어준뒤 small을 pop한 값을 large에 넣어준다. large의 개수가 많을때는(항상 small에 넣었다 빼서 large에 넣기 때문에 large가 항상 많거나 small과 같은 개수의 원소를 가짐) large에 넣었다 빼서 small에 넣어준다.

2. Median은 k가 홀수일때는 large의 첫번째 값이고, k가 짝수일때는 large의 첫번째 값과 small의 첫번째 값의 중간이다.

3. to_remove = defaultdict(int)를 통해 지워야 하는 숫자를 추적한다. 우리는 lazy removal를 사용할 것이므로, 오직 두개의 힙에서 첫번째 값이 지워야하는 값과 같을때만 지울것이다.

4. k다음부터의 원소를 힙들에 넣을때는, 이제는 항상 large에 넣었다 빼서 small에 넣어준다. 그리고 만약 지워야하는 숫자가 small의 첫번째 원소보다 크다면, 우리가 지워야 하는 값은 large에 있으므로 small에서 값을 뺀뒤 large에 넣어준다.

5. small과 large에 첫번째 원소가 지워야 하는 숫자이면 지워준다. 그리고 median값을 구해서 답에 넣는다.

Code

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        if not nums or not k:
            return []

        def get_median(small, large):
            return large[0] * 1. if k % 2 == 1 else (large[0]-small[0]) / 2.0

        small = [] # max heap
        large = [] # min heap
        for i in range(k):
            if len(small) == len(large):
                heapq.heappush(large, -heapq.heappushpop(small, -nums[i]))
            else:
                heapq.heappush(small, -heapq.heappushpop(large, nums[i]))

        ans = []
        ans.append(get_median(small, large))

        to_remove = defaultdict(int)
        for i in range(k, len(nums)):
            heapq.heappush(small, -heapq.heappushpop(large, nums[i]))

            num_to_delete = nums[i-k]
            if num_to_delete > -small[0]:
                heapq.heappush(large, -heapq.heappop(small))
            to_remove[num_to_delete] += 1

            while small and to_remove[-small[0]]:
                to_remove[-small[0]] -= 1
                heapq.heappop(small)
            while large and to_remove[large[0]]:
                to_remove[large[0]] -= 1
                heapq.heappop(large)

            ans.append(get_median(small, large))

        return ans

Complexity

Time Complexity: O()

Space Complexity: O()