295. Find Median from Data Stream

[zeikar]

Problem link

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

Problem Summary

데이터 스트림에서 현재까지의 중앙값을 출력하는 문제.

Solution

이것도 꽤 유명한 문제라 딱 보고 우선순위 큐 2개로 푸는 것이 생각나긴 했다.

간단히 설명하자면 현재까지의 데이터를 두 덩어리로 나눈다. left, right로 나누고 left는 max heap, right는 min heap이다. 이를 이용하면 중앙값은 left, right 에서 얻어오는 값으로 바로 구할 수 있다.
left, right 사이즈가 같다면 두 우선순위 큐의 최우선 값의 평균, 한쪽의 사이즈가 크다면 그 큐의 최우선 값이 중앙값이 된다.
삽입할 때 left, right의 최우선 값과 비교해서 큐에 집어넣고 밸런스를 맞추는 과정도 추가하였다.

코드에서는 파이썬 PriorityQueue 클래스를 사용했는데 여기서는 값을 지우지 않고 가져오는 메소드가 없어서... get 후 put 을 하였다.
그래서 성능이 생각보다 매우 안 좋은 듯.
heapq 클래스를 사용하면 더 빠르게 할 수 있다.

Source Code

from queue import PriorityQueue


class MedianFinder:

    def __init__(self):
        self.leftPriorityQueue = PriorityQueue()
        self.rightPriorityQueue = PriorityQueue()

        self.leftPriorityQueue.put(1000000)
        self.rightPriorityQueue.put(1000000)

    def addNum(self, num: int) -> None:
        left = -self.leftPriorityQueue.get()
        right = self.rightPriorityQueue.get()
        self.leftPriorityQueue.put(-left)
        self.rightPriorityQueue.put(right)

        if num < left:
            self.leftPriorityQueue.put(-num)
        else:
            self.rightPriorityQueue.put(num)

        if self.leftPriorityQueue.qsize() > self.rightPriorityQueue.qsize() + 1:
            val = -self.leftPriorityQueue.get()
            self.rightPriorityQueue.put(val)
        elif self.leftPriorityQueue.qsize() < self.rightPriorityQueue.qsize():
            val = self.rightPriorityQueue.get()
            self.leftPriorityQueue.put(-val)

    def findMedian(self) -> float:
        left = -self.leftPriorityQueue.get()
        right = self.rightPriorityQueue.get()
        self.leftPriorityQueue.put(-left)
        self.rightPriorityQueue.put(right)

        if self.leftPriorityQueue.qsize() == self.rightPriorityQueue.qsize():
            return (left + right) / 2
        elif self.leftPriorityQueue.qsize() > self.rightPriorityQueue.qsize():
            return left
        else:
            return right

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