Leetcode 295. Find Median (heapq 힙큐로 중간값 구하기)

https://leetcode.com/problems/find-median-from-data-stream/

Find Median from Data Stream - LeetCode

 

---

 

heapq는 잊을 만할때쯤 나오고 , O(NlogN)의 time complexity를 가지기 때문에 정렬에 종종 활용되는 것 같습니다.

이번 문제는 addNum의 기능에는 숫자를 추가하고 findMedian기능에서 중간값을 반환하도록 하는 문제인데요

 

가운데 선이 중간값 median이라면

median 보다 작은 값들 ( lower )

median보다 큰 값들 (higher)

이런식으로 중간값 기준 양쪽의 균형을 잡아가면서 비교해주고 중간값을 찾을 때, 각 리스트의 길이를 기반으로 lower의 가장 큰 값, higher의 가장 작은값을 기준으로 중간값을 계산해서 호출해주는 문제입니다.

 

*heapq는 가장 최소값을 기본적으로 반환하기 때문에 lower에서 higher로 값을 넘겨줄때, 그리고 lower에 값을 추가해줄 때 반드시 -1을 통해서 크기가 큰 값이 가장 큰 음수로 자리 잡아서 반환되도록 해야 합니다.

 

import heapq

class MedianFinder:
    def __init__(self):
        self.lower = []
        self.higher = []

    def addNum(self, num: int) -> None:
        heapq.heappush(self.lower,-heappushpop(self.higher,num))
        if len(self.lower) > len(self.higher):
            heapq.heappush(self.higher,-(heapq.heappop(self.lower)))

    def findMedian(self) -> float:         
        if len(self.lower) < len(self.higher):
                return self.higher[0]
        return (self.lower[0]*-1 + self.higher[0])/2


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()

 

여러 풀이를 보면서 답안을 구성했는데요. lower과 higher이라는 리스트를 각각 만들어주고, lower가 higher보다 길면 higher로 값을 넘겨주며 균형을 맞춰가다가, 중간값을 계산해야 할 때 전체 길이가 짝수가 아닌 경우에 higher의 가장 작은 값을 반환하도록 했습니다.

 

---

 

다른 코드들을 보면서 분석해보겠습니당

 

간결한 풀이)

from heapq import *

class MedianFinder:

    def __init__(self):
        self.heaps = [], []

    def addNum(self, num):
        small, large = self.heaps
        heappush(small, -heappushpop(large, num))
        if len(large) < len(small):
            heappush(large, -heappop(small))

    def findMedian(self):
        small, large = self.heaps
        if len(large) > len(small):
            return float(large[0])
        return (large[0] - small[0]) / 2.0

 

저처럼 self.lower, self.higher을 매번 쓴게 아니라 heaps라는 값을 초기에 설정해주고 small, large로 각 함수에서

나누어지도록 해서 코드가 깔끔하게 되도록 했습니다.

 

마지막 median을 찾을 때도 길이가 다른 예외의 경우를 if 문에 먼저 제시를 해서 if 문의 불필요한 반복을 막았습니다.

 

풀이2 )

from heapq import *

class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.min_heap: List[int] = []
        self.max_heap: List[int] = []
        

    def addNum(self, num: int) -> None:
        heappush(self.max_heap, -heappushpop(self.min_heap, num))
        
        if len(self.max_heap) > len(self.min_heap):
            heappush(self.min_heap, -heappop(self.max_heap))
        

    def findMedian(self) -> float:
        has_even_count = len(self.max_heap) == len(self.min_heap)

        if has_even_count:
            return (-self.max_heap[0] + self.min_heap[0]) / 2.0
        return float(self.min_heap[0])

 

이 코드 역시 방향성은 같지만 median보다 큰 값에 값을 먼저 넣냐, median 기준 더 작은 리스트에 값을 먼저 넣냐의 차이가 있습니다. 마지막에 has_even_count로 짝수개의 값을 가지는지 여부를 boolean으로 판별하도록 했습니다.