LeetCode: 295. Find Median from Data Stream

면접 형식의 문제여서 문제를 맞추는 것은 어렵지 않다.

다만 어느정도로 효율적인 알고리즘(시간 및 메모리 측면)을 구현하는지가 관건이다.

 

가장 간단한 방법은 모든 숫자를 저장하고,

중앙값 search시마다 정렬하는 방법이다.

add: O(1)
find: O(nlogn)

위의 이유로 비효율적이다.

class MedianFinder:

    def __init__(self):
        self.arr = []
        self.count = 0

    def addNum(self, num: int) -> None:
        self.arr.append(num)
        self.count += 1

    def findMedian(self) -> float:
        self.arr.sort()
        if self.count % 2 == 0:
            return (self.arr[self.count//2-1]+self.arr[self.count//2])/2
        else:
            return self.arr[self.count//2]

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

 

heap 2개를 사용하면 문제를 효율적으로 풀 수 있다.

add: O(logn)

find: O(1)

import heapq
class MedianFinder:
    def __init__(self):
        # self.h1: max heap of small half
        # self.h2: min heap of the other half
        self.h1, self.h2 = [], []

    def addNum(self, num: int) -> None:
        if len(self.h1) == len(self.h2):
            heapq.heappush(self.h2, -heapq.heappushpop(self.h1, -num))
        else:
            heapq.heappush(self.h1, -heapq.heappushpop(self.h2, num))
        
    def findMedian(self) -> float:
        if len(self.h1) == len(self.h2):
            return (-self.h1[0] + self.h2[0])/2
        return self.h2[0]
        

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