끊임없이 새로운 숫자가 유입되는 데이터 스트림이 있다고 가정해 봅시다. 이 스트림에서 지금까지 들어온 모든 숫자의 중앙값(median)을 빠르게 찾아내는 시스템을 만들어야 합니다.
중앙값은 정렬된 리스트의 가운데 값입니다. 리스트 길이가 홀수라면 정확히 가운데 원소 하나를 그대로 반환하면 되고, 짝수라면 가운데 두 원소의 평균을 계산하면 됩니다.
이 문제를 해결하기 위해 두 개의 메서드를 구현해야 합니다.
addNum()— 스트림에 숫자를 추가하는 메서드findMedian()— 지금까지 추가된 모든 숫자의 중앙값을 반환하는 메서드
해결 방법: 두 개의 힙(Heap) 활용
가장 효율적인 접근 방식은 최대 힙(max-heap)과 최소 힙(min-heap)을 함께 사용하는 것입니다. 왼쪽 힙은 작은 절반을 내림차순으로 관리하고, 오른쪽 힙은 큰 절반을 오름차순으로 관리하여 두 힙의 균형을 유지하면 중앙값을 O(log n) 시간에 구할 수 있습니다.
알고리즘 단계
- 우선순위 큐
left(최대 힙)와right(최소 힙)를 선언합니다.
addNum(num) 메서드
left가 비어 있거나num이left의 최상단(top) 원소보다 작으면 →num을left에 삽입합니다.- 그렇지 않으면 →
num을right에 삽입합니다. left의 크기가right보다 작으면:temp := right.top()right에서 최상단 원소를 제거한 뒤temp를left에 삽입합니다.
left의 크기에서right의 크기를 뺀 값이 1보다 크면:temp := left.top()left에서 최상단 원소를 제거한 뒤temp를right에 삽입합니다.
findMedian() 메서드
left의 크기가right보다 크면left.top()을 반환하고, 그렇지 않으면(left.top() + right.top()) / 2를 반환합니다.
C++ 구현 예제
아래 코드를 통해 동작 방식을 더 쉽게 이해할 수 있습니다.
#include <bits/stdc++.h>
using namespace std;
typedef double lli;
class MedianFinder {
priority_queue <int> left;
priority_queue <int, vector <int>, greater<int>> right;
public:
void addNum(int num) {
if(left.empty() || num<left.top()){
left.push(num);
}else right.push(num);
if(left.size()<right.size()){
lli temp = right.top();
right.pop();
left.push(temp);
}
if(left.size()-right.size()>1){
lli temp = left.top();
left.pop();
right.push(temp);
}
}
double findMedian() {
return
left.size()>right.size()?left.top():(left.top()+right.top())*0.5;
}
};
main(){
MedianFinder ob;
ob.addNum(10);
ob.addNum(15);
cout << ob.findMedian() << endl;
ob.addNum(25);
ob.addNum(30);
cout << ob.findMedian() << endl;
ob.addNum(40);
cout << ob.findMedian();
}입력
addNum(10); addNum(15); findMedian(); addNum(25); addNum(30); findMedian(); addNum(40); findMedian();
출력
12.5 20 25
동작 과정 살펴보기
- 10, 15 추가 후: 정렬된 상태는 [10, 15]이며, 두 수의 평균인 12.5가 출력됩니다.
- 25, 30 추가 후: 스트림은 [10, 15, 25, 30]이 되고, 가운데 두 수 15와 25의 평균인 20이 출력됩니다.
- 40 추가 후: 스트림은 [10, 15, 25, 30, 40]이 되어 홀수 개이므로 정확히 가운데 값인 25가 출력됩니다.
이처럼 두 개의 우선순위 큐를 사용하면 매번 전체 데이터를 다시 정렬하지 않고도 스트림의 중앙값을 효율적으로 추적할 수 있습니다.