C++로 실시간 숫자 스트림의 중앙값 구하기 – 힙을 활용한 효율적인 알고리즘

이 문제에서는 정수를 지속적으로 읽어들이는 데이터 스트림이 주어집니다. 우리의 과제는 스트림에서 요소를 하나씩 읽을 때마다, 그 시점까지 입력된 모든 요소의 중앙값(median)을 계산하는 프로그램을 작성하는 것입니다.

중앙값(Median)이란 정렬된 수열(오름차순 또는 내림차순)에서 가운데에 위치한 원소를 의미합니다.

중앙값 계산 방법

원소의 개수가 홀수일 때는 가운데 원소 하나가 중앙값이 됩니다.

원소의 개수가 짝수일 때는 가운데 두 원소의 평균이 중앙값이 됩니다.

예시로 이해하기

입력 − 3, 65, 12, 20, 1

각 입력이 들어올 때마다 중앙값은 다음과 같이 변화합니다.

입력 3 : 수열 (3) → 중앙값 3
입력 65 : 수열 (3, 65) → 중앙값 34
입력 12 : 수열 (3, 12, 65) → 중앙값 12
입력 20 : 수열 (3, 12, 20, 65) → 중앙값 16
입력 1 : 수열 (1, 3, 12, 20, 65) → 중앙값 12

위 예시에서는 중앙값 계산을 쉽게 하기 위해 정렬된 수열을 사용했습니다.

문제 해결 접근 방법

이 문제를 해결하는 방법은 여러 가지가 있습니다. 매 단계마다 데이터를 정렬한 뒤 중앙값을 찾는 방법, 자가 균형 이진 탐색 트리(self-balancing BST)를 사용하는 방법, 그리고 힙(heap)을 사용하는 방법이 대표적입니다. 그중에서도 힙을 활용한 방법이 가장 효율적이고 유망한 해결책으로 꼽힙니다.

핵심 아이디어는 다음과 같습니다. 최대 힙(max-heap)에는 현재까지 입력된 값 중 작은 절반을, 최소 힙(min-heap)에는 큰 절반을 저장합니다. 두 힙의 크기 차이가 1을 넘지 않도록 균형을 유지하면, 중앙값은 항상 두 힙의 루트(top)에서 상수 시간(O(1))에 구할 수 있고, 새 원소를 삽입할 때도 로그 시간(O(log n))이면 충분합니다. 즉, 매번 전체 데이터를 정렬하는 O(n log n) 방식보다 훨씬 빠릅니다.

구현 예제

다음은 위에서 설명한 알고리즘을 C++로 구현한 프로그램입니다.

#include <iostream>
using namespace std;
#define MAX_HEAP_SIZE (128)
#define ARRAY_SIZE(a) sizeof(a)/sizeof(a[0])
void swap(int &a, int &b){
    int temp = a;
    a = b;
    b = temp;
}
bool Greater(int a, int b){
    return a > b;
}
bool Smaller(int a, int b){
    return a < b;
}
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 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]) ) {
                swap(A[i], A[p]);
                heapify(p);
            }
        }
        int deleteTop(){
            int del = -1;
            if( heapSize > -1){
                del = A[0];
                swap(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 {
    private:
    public:
        MaxHeap() : Heap(new int[MAX_HEAP_SIZE], &Greater) { }
        int GetTop() {
            return top();
        }
        int ExtractTop() {
            return deleteTop();
        }
        int GetCount() {
            return count();
        }
        bool Insert(int key) {
            return insertHelper(key);
        }
};
class MinHeap : public Heap{
    private:
    public:
        MinHeap() : Heap(new int[MAX_HEAP_SIZE], &Smaller) { }
        int GetTop() {
            return top();
        }
        int ExtractTop() {
            return deleteTop();
        }
        int GetCount() {
            return count();
        }
        bool Insert(int key) {
            return insertHelper(key);
        }
};
int findMedian(int e, int &median, Heap &left, Heap &right){
    switch(Signum(left.GetCount(), right.GetCount())){
        case 0: if( e < median ) {
            left.Insert(e);
            median = left.GetTop();
        }
        else{
            right.Insert(e);
            median = right.GetTop();
        }
        break;
        case 1: if( e < median ){
            right.Insert(left.ExtractTop());
            left.Insert(e);
        }
        else
            right.Insert(e);
            median = ((left.GetTop()+right.GetTop())/2);
        break;
        case -1: if( e < median )
            left.Insert(e);
        else {
            left.Insert(right.ExtractTop());
            right.Insert(e);
        }
        median = ((left.GetTop()+right.GetTop())/2);
        break;
    }
    return median;
}
void printMedianStream(int A[], int size){
    int median = 0;
    Heap *left = new MaxHeap();
    Heap *right = new MinHeap();
    for(int i = 0; i < size; i++) {
        median = findMedian(A[i], median, *left, *right);
        cout<<"Median of elements : ";
        for(int j = 0; j<=i;j++) cout<<A[j]<<" ";
        cout<<"is "<<median<<endl;
    }
}
int main(){
    int A[] = {12, 54, 9, 6, 1};
    int size = ARRAY_SIZE(A);
    printMedianStream(A, size);
    return 0;
}

출력 결과

Median of elements : 12 is 12
Median of elements : 12 54 is 33
Median of elements : 12 54 9 is 12
Median of elements : 12 54 9 6 is 10
Median of elements : 12 54 9 6 1 is 9

출력 결과를 보면, 새로운 숫자가 스트림에 들어올 때마다 그 시점까지의 전체 수열에 대한 중앙값이 즉시 갱신되는 것을 확인할 수 있습니다. 예를 들어 {12, 54}가 입력되었을 때는 두 값의 평균인 33이, 다섯 번째 값 1까지 입력되었을 때는 정렬된 수열 {1, 6, 9, 12, 54}의 가운데 값인 9가 중앙값으로 출력됩니다.

이처럼 최대 힙과 최소 힙을 조합하면 스트림 데이터가 아무리 길어져도 중앙값을 실시간으로 효율적으로 추적할 수 있으며, 이 기법은 통계 분석, 실시간 모니터링, 순위 집계 등 다양한 분야에서 활용됩니다.