295. Find Median from Data Stream

Find Median from Data Stream - LeetCode

정수가 입력되고 입력된 정수는 어떤 배열에 오름차순으로 정렬되어 저장된다고 가정하자.

이때 정렬된 배열의 중간값(median)을 반환하는 문제다.

 

예시

Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Output
[null, null, null, 1.5, null, 2.0]

Explanation
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

 

접근법

정수를 입력할 때마다 정렬을 한다면 O(nlogn)이 findMedian()을 호출할 때마다 발생할 것이다. 그렇기 때문에 시간 복잡도를 줄이면서 요구하는 값을 반환해야 한다.

 

약간 Greedy 하게 생각해보면 다음과 같다.

• 우리가 원하는 값은 정렬된 배열의 중간값이다.
• 만일 정렬된 배열을 반으로 쪼갠 배열 2개를 a, b라고 한다면 우리가 원하는 값은 배열 a의 마지막 값, 혹은 배열 b의 첫 번째 값, 혹은 이 둘을 더하고 2로 나눈 값이다.
• 위에서 말했던 배열 a, b를 유지한다면 findMedian()을 O(1)로 수행할 수 있을 것이다.

 

1. 우선순위 큐 x 2

from heapq import *


class MedianFinder:

    def __init__(self):
        self.max_heap = []
        self.min_heap = []

    def addNum(self, num: int) -> None:
        if not self.max_heap:
            heappush(self.max_heap, -num)
            return

        if self.min_heap and num > self.min_heap[0]:
            heappush(self.min_heap, num)
        else:
            heappush(self.max_heap, -num)
        if len(self.max_heap) - len(self.min_heap) > 1:
            heappush(self.min_heap, -heappop(self.max_heap))
        elif len(self.min_heap) - len(self.max_heap) > 1:
            heappush(self.max_heap, -heappop(self.min_heap))

    def findMedian(self) -> float:
        if len(self.max_heap) != len(self.min_heap):
            return -self.max_heap[0] if len(self.max_heap) > len(self.min_heap) else self.min_heap[0]
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2

우선순위 큐 2개를 이용하면 해결할 수 있다.

max_heap은 위에서 말했던 정렬된 배열을 쪼갠 왼쪽 배열이고, min_heap은 오른쪽 배열이다.

 

max_heap에 우선적으로 num을 넣어주고, 입력될 num이 min_heap [0]보다 크다면 min_heap에 넣어주고 그렇지 않다면 max_heap에 넣어준다.

 

그리고 max_heap과 min_heap의 길이 차이가 2 이상 난다면 길이가 긴 쪽의 heap에서 짧은 쪽의 heap으로 값을 넣어준다.

 

이 상태를 유지하면 문제에서 요구하는 중간값을 정렬 없이 찾을 수 있다. 삽입할 때 O(logn)이 소요되지만 O(nlogn)보단 훨씬 작은 값이기 때문에 TLE 없이 문제를 해결할 수 있다.

 

1655번: 가운데를 말해요

이 문제와 똑같은 문제를 백준에서 풀어본 경험이 있어서 풀 수 있었지, 아니었으면 사실 접근도 못했을 문제긴 하다;