C++로 두 개의 힙을 활용해 실시간으로 중앙값 찾기

MedianClass라는 클래스를 구현해야 한다고 가정해 봅시다. 이 클래스에는 다음과 같은 메서드가 포함되어야 합니다.

add(value): 데이터 구조에 새로운 값을 추가합니다.
median(): 현재 데이터 구조에 저장된 모든 숫자의 중앙값을 계산하여 반환합니다.

예를 들어 5, 3, 8을 차례로 추가한 후 중앙값을 조회하면 결과는 5.0입니다. 이어서 9를 추가한 뒤 다시 중앙값을 조회하면 결과는 6.5가 됩니다.

해결 접근 방식

이 문제는 최대 힙(max-heap) 하나와 최소 힙(min-heap) 하나, 총 두 개의 우선순위 큐를 사용하면 매우 효율적으로 해결할 수 있습니다. 왼쪽 힙(left)은 전체 데이터 중 작은 절반을, 오른쪽 힙(right)은 큰 절반을 관리하므로, 새 값이 추가될 때마다 두 힙의 균형만 유지해 주면 중앙값을 항상 빠르게 얻을 수 있습니다.

구체적인 알고리즘은 다음과 같습니다.

우선순위 큐 left와 right를 정의합니다.
addNum 메서드: 숫자를 입력받아 다음 과정을 수행합니다.
만약 left가 비어 있거나, 입력된 num이 left의 최상단(top) 원소보다 작다면:
- num을 left에 삽입합니다.
그렇지 않다면:
- num을 right에 삽입합니다.
만약 left의 크기가 right의 크기보다 작다면:
- temp := right의 최상단 원소
- right에서 해당 원소를 삭제합니다.
- temp를 left에 삽입합니다.
만약 (left의 크기 − right의 크기)가 1보다 크다면:
- temp := left의 최상단 원소
- left에서 해당 원소를 삭제합니다.
- temp를 right에 삽입합니다.
findMedian 메서드: 다음과 같이 동작합니다.

left의 크기가 right의 크기보다 크면 left의 최상단 원소를 그대로 반환하고, 그렇지 않으면 (left의 최상단 원소 + right의 최상단 원소) / 2를 반환합니다.

이 방식의 시간 복잡도는 값 추가 시 O(log n), 중앙값 조회 시 O(1)로, 데이터가 스트림처럼 지속적으로 들어오는 상황에서 특히 유용합니다. 아래 예제 코드를 통해 더 자세히 살펴보겠습니다.

예제 코드

#include <bits/stdc++.h>
using namespace std;
typedef double lli;
class MedianClass {
    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(){
    MedianClass ob;
    ob.addNum(5);
    ob.addNum(3);
    ob.addNum(8);
    cout << ob.findMedian() << " ";
    ob.addNum(9);
    cout << ob.findMedian() << endl;
}

입력

ob.addNum(5);
ob.addNum(3);
ob.addNum(8);
cout << ob.findMedian() << endl;
ob.addNum(9);
cout << ob.findMedian() << endl;

출력

5.0
6.5

동작 원리 요약

첫 번째 단계에서 5, 3, 8이 추가되면 왼쪽 힙에는 {3, 5}, 오른쪽 힙에는 {8}이 저장되어 두 힙의 크기가 같으므로 중앙값은 (5 + 8) / 2 = 6.5가 아니라, 실제 코드 실행 시 힙의 균형 조건에 따라 5.0이 출력됩니다. 이후 9가 추가되면 오른쪽 힙은 {8, 9}가 되고, 두 힙의 최상단 값인 5와 8의 평균인 6.5가 중앙값으로 반환됩니다. 이처럼 두 개의 힙을 활용하면 데이터가 계속 늘어나더라도 항상 효율적으로 중앙값을 유지할 수 있습니다.