C++で整数ストリームの中央値を効率的に求める方法（ヒープを使ったアルゴリズム解説）

問題概要

データストリームから順次整数が読み込まれる状況を考えます。このとき、それまでに読み込んだ要素の中から中央値（メディアン）を、効率的な方法で求めることが課題です。

処理の流れは以下のようになります。

• 1つ目の要素を読み込んだ後：10 → 中央値は 10
• 2つ目の要素を読み込んだ後：10, 20 → 中央値は 15（平均値）
• 3つ目の要素を読み込んだ後：10, 20, 30 → 中央値は 20

このように、要素が追加されるたびに中央値を更新し続ける必要があります。

アルゴリズムの考え方

この問題を効率的に解くには、最大ヒープ（max heap）と最小ヒープ（min heap）を組み合わせる手法が有効です。手順は以下の通りです。

1. 左右に分割したヒープ構成： 有効な中央値より小さい要素群を格納するために左側に最大ヒープを、大きい要素群を格納するために右側に最小ヒープを使用します。
2. 要素数のバランス維持： 新しい要素を処理した後、両ヒープの要素数の差は常に1個以内に保ちます。
3. 要素数が等しい場合： 両ヒープのルート（先頭）データの平均値を有効な中央値として採用します。
4. 要素数が不均衡な場合： 要素数が多い方のヒープのルート値をそのまま有効な中央値として選択します。

この手法により、各要素の挿入と中央値の取得を O(log n) の計算量で実現でき、全体的にも非常に効率的な処理が可能になります。

C++による実装例

#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{
private:
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{
private:
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:\n";
    printMedian(A, size);
    return 0;
}

実装のポイント

• Heap クラスは抽象基底クラスとして設計され、比較関数をコンストラクタで受け取ることで、最大ヒープと最小ヒープのコードを共通化しています。
• Signum 関数によって左右のヒープの要素数の関係（多い・等しい・少ない）を判定し、それぞれの場合に応じて適切に再バランスを行っています。
• 新しい要素が現在の中央値より小さい場合は左側（最大ヒープ）へ、そうでなければ右側（最小ヒープ）へ振り分けることで、常に「左ヒープ ≤ 中央値 ≤ 右ヒープ」という不変条件を維持します。

実行結果

上記のプログラムをコンパイルして実行すると、以下の出力が得られます。

Result:
10
15
20

この結果から、要素を読み込むたびに中央値が正しく更新されていることが確認できます。10 を読み込んだ時点では中央値は 10、20 を追加すると平均の 15、さらに 30 を追加すると 20 となります。