문제 정의
데이터 스트림을 통해 정수가 하나씩 계속 입력되는 상황을 가정해 보겠습니다. 이때 지금까지 읽어 들인 모든 원소의 중앙값(median)을 매 순간 효율적으로 구해야 합니다.
입력이 10, 20, 30 순서로 들어온다면 중앙값은 다음과 같이 변합니다.
- 첫 번째 원소 10을 읽은 후 → {10}의 중앙값은 10
- 두 번째 원소 20을 읽은 후 → {10, 20}의 중앙값은 15
- 세 번째 원소 30을 읽은 후 → {10, 20, 30}의 중앙값은 20
접근 방법: 두 개의 힙(Heap) 활용
새 원소가 들어올 때마다 전체 데이터를 다시 정렬하는 것은 비효율적입니다. 대신 최대 힙(max heap)과 최소 힙(min heap)을 함께 사용하면 각 원소를 O(log n) 시간에 처리하면서 중앙값을 O(1)에 얻을 수 있습니다.
- 힙 구성: 왼쪽에는 최대 힙을 두어 유효 중앙값 이하의 원소들을 저장하고, 오른쪽에는 최소 힙을 두어 유효 중앙값보다 큰 원소들을 저장합니다.
- 균형 유지: 새 원소를 처리한 뒤에는 두 힙에 담긴 원소 개수의 차이가 항상 1 이하가 되도록 유지합니다.
- 개수가 같을 때: 두 힙의 루트(root) 값의 평균을 유효 중앙값으로 사용합니다.
- 균형이 깨졌을 때: 원소 수가 더 많은 쪽 힙의 루트 값을 그대로 중앙값으로 선택합니다.
시간 복잡도
원소 삽입 및 힙 재조정: O(log n)
중앙값 조회: O(1)
C++ 구현 예제
아래 코드는 추상 클래스 Heap을 기반으로 MaxHeap과 MinHeap을 정의하고, 스트림의 각 원소를 처리할 때마다 현재까지의 중앙값을 출력합니다. 핵심 함수인 getMedian은 Signum으로 두 힙의 크기 관계(양수·0·음수)를 판별한 뒤 세 가지 경우로 나누어 처리합니다.
#include <iostream>
using namespace std;
#define MAX_HEAP_SIZE (128)
#define ARRAY_SIZE(a) sizeof(a)/sizeof(a[0])
inline void Exch(int &a, int &b){
int aux = a;
a = b;
b = aux;
}
bool Greater(int a, int b){
return a > b;
}
bool Smaller(int a, int b){
return a < b;
}
int Average(int a, int b){
return (a + b) / 2;
}
int Signum(int a, int b){
if(a == b){
return 0;
}
return a < b ? -1 : 1;
}
class Heap{
public:
Heap(int *b, bool (*c)(int, int)) : A(b), comp(c){
heapSize = -1;
}
virtual ~Heap(){
if(A){
delete[] A;
}
}
virtual bool Insert(int e) = 0;
virtual int GetTop() = 0;
virtual int ExtractTop() = 0;
virtual int GetCount() = 0;
protected:
int left(int i){
return 2 * i + 1;
}
int right(int i){
return 2 * (i + 1);
}
int parent(int i){
if(i <= 0){
return -1;
}
return (i - 1) / 2;
}
int *A;
bool (*comp)(int, int);
int heapSize;
int top(void){
int max = -1;
if(heapSize >= 0){
max = A[0];
}
return max;
}
int count(){
return heapSize + 1;
}
void heapify(int i){
int p = parent(i);
if(p >= 0 && comp(A[i], A[p])){
Exch(A[i], A[p]);
heapify(p);
}
}
int deleteTop(){
int del = -1;
if(heapSize > -1){
del = A[0];
Exch(A[0], A[heapSize]);
heapSize--;
heapify(parent(heapSize + 1));
}
return del;
}
bool insertHelper(int key){
bool ret = false;
if(heapSize < MAX_HEAP_SIZE){
ret = true;
heapSize++;
A[heapSize] = key;
heapify(heapSize);
}
return ret;
}
};
class MaxHeap : public Heap{
public:
MaxHeap() : Heap(new int[MAX_HEAP_SIZE], &Greater){ }
~MaxHeap(){ }
int GetTop(){ return top(); }
int ExtractTop(){ return deleteTop(); }
int GetCount(){ return count(); }
bool Insert(int key){ return insertHelper(key); }
};
class MinHeap : public Heap{
public:
MinHeap() : Heap(new int[MAX_HEAP_SIZE], &Smaller){ }
~MinHeap(){ }
int GetTop(){ return top(); }
int ExtractTop(){ return deleteTop(); }
int GetCount(){ return count(); }
bool Insert(int key){ return insertHelper(key); }
};
// 새 원소를 처리하며 유효 중앙값을 갱신하는 함수
int getMedian(int e, int &m, Heap &l, Heap &r){
int sig = Signum(l.GetCount(), r.GetCount());
switch(sig){
case 1: // 왼쪽 힙이 더 많은 경우
if(e < m){
r.Insert(l.ExtractTop());
l.Insert(e);
} else {
r.Insert(e);
}
m = Average(l.GetTop(), r.GetTop());
break;
case 0: // 두 힙의 크기가 같은 경우
if(e < m){
l.Insert(e);
m = l.GetTop();
} else {
r.Insert(e);
m = r.GetTop();
}
break;
case -1: // 오른쪽 힙이 더 많은 경우
if(e < m){
l.Insert(e);
} else {
l.Insert(r.ExtractTop());
r.Insert(e);
}
m = Average(l.GetTop(), r.GetTop());
break;
}
return m;
}
void printMedian(int A[], int size){
int m = 0;
Heap *left = new MaxHeap();
Heap *right = new MinHeap();
for(int i = 0; i < size; ++i){
m = getMedian(A[i], m, *left, *right);
cout << m << endl;
}
delete left;
delete right;
}
// 드라이버 코드
int main(){
int A[] = {10, 20, 30};
int size = ARRAY_SIZE(A);
cout << "Result:" << endl;
printMedian(A, size);
return 0;
}실행 결과
위 프로그램을 컴파일하여 실행하면 다음과 같은 결과가 출력됩니다.
Result: 10 15 20
마무리
이처럼 최대 힙과 최소 힙을 균형 있게 운용하면 데이터를 매번 정렬하지 않고도 스트림으로 흘러들어오는 정수의 실시간 중앙값을 효율적으로 추적할 수 있습니다. 왼쪽 힙은 중앙값 이하의 절반을, 오른쪽 힙은 중앙값 초과의 절반을 담당하므로, 중앙값은 항상 두 힙의 루트에서 바로 얻을 수 있다는 점이 이 알고리즘의 핵심입니다.