295. Find Median from Data Stream
[Find Median from Data Stream - LeetCode
Level up your coding skills and quickly land a job. This is the best place to expand your knowledge and get prepared for your next interview.
leetcode.com](https://leetcode.com/problems/find-median-from-data-stream/)
정수가 입력되고 입력된 정수는 어떤 배열에 오름차순으로 정렬되어 저장된다고 가정하자.
이때 정렬된 배열의 중간값(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번: 가운데를 말해요
첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -1
www.acmicpc.net](https://www.acmicpc.net/problem/1655)
이 문제와 똑같은 문제를 백준에서 풀어본 경험이 있어서 풀 수 있었지, 아니었으면 사실 접근도 못했을 문제긴 하다;