[Two Heaps] Heap-PrioiryQueue: 295. Find Median from Data Stream ★★★

295. Find Median from Data Stream

 

https://youtu.be/itmhHWaHupI?si=QWJJw3c9WUZ0gmFN

 

class MedianFinder:
    # 힙에서 데이터 추가/삭제는 O(log N) 이고, 검색은 O(1) 임
    # small heap 은 max heap 으로, large heap 은 min heap 으로..
    # max heap 을 pop 하면, 
    #             가장 큰값이 추출되므로 이것을 large heap으로 바로 보낼 수 있음
    # 우선 small heap 에 데이터 넣고
    # 발란스가 small heap 이 더 크면, small heap 의 최대갑을 large heap 으로 넣고
    # 여기서 발란스는 
    #   두개 길이 차이가 <=1 이내이면, 그냥 두고.
    #       길이 차이가 >=2 이상이면, 발란스 맞추기
    # 그리고, small heap 의 최대값이 large heap 의 최소값보다 크면,
    #       small heap 의 최대값을 빼서 large heap 에 넣기.
    # median 값은 small heap 의 최대값과 large heap 의 최소값을 평균내면 됨


    def __init__(self):
        # two heaps, large (min heap), small (max heap)
        # heaps should be equal size
        self.small = []  # maxheap
        self.large = []  # minheap

    def addNum(self, num: int) -> None:
        # 모든 시간 연산은 O (log n)
        # max heap을 위해 -1 곱하기, python 은 min heap 만 지원        
        # heapq.heappush(self.small, -1 * num) 
        
        # small 은 maxheap 으로
        self.small.append(-num)
        heapq.heapify(self.small)
        heapq.heapify(self.large)

		"""
        # 둘다 있는데, 
        # small 부분의 최대값이 large 부분의 최소값보다 크면
        # small 최대값을 빼서, large 로 넣어줘야함
        """
        # make sure every num small is <= every num in large
        # (-1 * self.small[0]) 에서 가장 큰값 
        # self.large[0] 에서 가장 작은 값
        if (self.small and self.large and 
            (-1 * self.small[0]) > self.large[0]):
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)

        # uneven size?
        """
        # small 이 크기가 더 크면, small 의 최대값이 -> large로
        """
        if len(self.small) > len(self.large) + 1:
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)
		
        """
        # large 이 크기가 더 크면, large 의 최소값이 -> small로
        """
        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -1 * val)
        

    def findMedian(self) -> float:
        # 모든 시간 연산은 O (1)
        # odd len
        if len(self.large) > len(self.small):
            return self.large[0]
        elif len(self.large) < len(self.small):
            return (-1)*self.small[0]
        else: # len(self.large) == len(self.small)
            return (self.large[0] + (-1)*self.small[0])/2
        


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

 

최종 코드

import heapq

class MedianFinder:

    def __init__(self):
        self.smallMaxHeap = []
        self.largeMinHeap = []

    def addNum(self, num: int) -> None:
        # 일단 모든 데이터를 small heap 에 모두 넣음
        # python 은 기본이 min heap 임
        # max heap 으로 만들어야 하기에. -1을 곱해줌        
        heapq.heappush(self.smallMaxHeap, -1*num)

        # 두개의 힙에 데이터가 있고, 두 힙의 정렬 순서가 깨지는 경우
        if self.smallMaxHeap and self.largeMinHeap and (-1*self.smallMaxHeap[0] > self.largeMinHeap[0]):
            val = -1*heapq.heappop(self.smallMaxHeap)
            heapq.heappush(self.largeMinHeap, val)

        # 두개 힙의 사이즈 정렬
        # small heap 이 더 많으면, small -> large
        if len(self.smallMaxHeap) > len(self.largeMinHeap) + 1: 
            val = -1*heapq.heappop(self.smallMaxHeap)
            heapq.heappush(self.largeMinHeap, val)

        # large heap 이 더 많으면, large -> small
        if len(self.smallMaxHeap) +1 < len(self.largeMinHeap): 
            val = heapq.heappop(self.largeMinHeap)
            heapq.heappush(self.smallMaxHeap, -1*val)

    def findMedian(self) -> float:
        # small heap 개수가 더 많으면, small heap 의 최대값이 중간값
        if len(self.smallMaxHeap) > len(self.largeMinHeap): 
            return -1*self.smallMaxHeap[0]
        elif len(self.smallMaxHeap) < len(self.largeMinHeap):
            return  self.largeMinHeap[0]
        else: # 개수가 같으면, 두개의 평균값이 중간값
            return 0.5*(-1*self.smallMaxHeap[0] + self.largeMinHeap[0] )


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