LeetCode/NeetCode

[Two Heaps] 480. Sliding Window Median

hyunkookim 2025. 4. 8. 11:53

480. Sliding Window Median

 

heapq을 사용해 Two Heaps 방식으로 구현한 것입니다.
앞서 봤던 방식과 비슷하지만,

좀 더 간결하게 balance 변수로 크기 조정을 처리한 형태입니다.

 

import heapq
from collections import defaultdict
from typing import List

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        small, large = [], []          # small: 최대 힙 (음수 저장), large: 최소 힙
        d = defaultdict(int)           # 삭제 예약 딕셔너리

        # --- 초기 윈도우 설정 ---
        for i in range(k):
            heapq.heappush(small, -nums[i])  # 최대 힙에 음수로 저장 (heapq는 최소 힙만 지원)
        for i in range(k // 2):
            # 균형을 위해 절반은 large로 이동
            heapq.heappush(large, -heapq.heappop(small))

        # --- 첫 번째 중간값 저장 ---
        # 홀수 → small의 루트가 중간값
        # 짝수 → small과 large의 루트 평균
        res = [-small[0] if k & 1 else (large[0] - small[0]) / 2]

        # --- 슬라이딩 윈도우 시작 ---
        for i in range(k, len(nums)):
            out_num = nums[i - k]      # 윈도우에서 빠지는 값
            in_num = nums[i]           # 윈도우에 새로 들어오는 값
            d[out_num] += 1            # 삭제 예약 (언젠가 heap의 top에서 제거됨)

            # balance: 두 힙 간 크기 차이 추적
            # out_num이 small에서 제거되면 small 크기 줄어듦 → balance -1
            # out_num이 large에서 제거되면 large 크기 줄어듦 → balance +1
            balance = -1 if small and out_num <= -small[0] else 1

            # 새로 들어온 값 in_num 삽입
            if small and in_num <= -small[0]:
                heapq.heappush(small, -in_num)  # 최대 힙
                balance += 1                   # small 크기 증가
            else:
                heapq.heappush(large, in_num)  # 최소 힙
                balance -= 1                   # large 크기 증가

            # balance 값으로 힙 균형 조정
            if balance > 0:
                # small이 커졌으니 하나 large로 이동
                heapq.heappush(large, -heapq.heappop(small))
            if balance < 0:
                # large가 커졌으니 하나 small로 이동
                heapq.heappush(small, -heapq.heappop(large))

            # lazy deletion: top에 삭제 예약된 값 있으면 제거
            while small and d[-small[0]] > 0:
                d[-heapq.heappop(small)] -= 1

            while large and d[large[0]] > 0:
                d[heapq.heappop(large)] -= 1

            # 현재 윈도우의 중간값 계산 후 저장
            res.append(-small[0] if k & 1 else (large[0] - small[0]) / 2)

        return res  # 모든 윈도우의 중간값 리스트 반환

 

k & 1이 왜 홀수인지 확인하는 조건이 되는지 설명드릴게요.


✅ k & 1 이란?

이건 비트 연산자 AND를 사용한 표현입니다.
정수 k의 **가장 마지막 비트(1의 자리)**가 1인지 확인하는 연산이에요.


✅ 이진수로 살펴보기

k 이진수 k & 1 결과 의미
3 011 011 & 001 → 001 1 → 홀수
4 100 100 & 001 → 000 0 → 짝수
5 101 101 & 001 → 001 1 → 홀수

즉, k & 1 == 1이면 홀수,
k & 1 == 0이면 짝수입니다.


✅ 왜 쓰는가?

if k & 1:
    # k가 홀수일 때
    median = -small[0]
else:
    # k가 짝수일 때
    median = (large[0] - small[0]) / 2
 
이런 식으로 홀/짝을 빠르게 판별할 수 있어서 사용합니다.