C++でデータストリームから中央値を効率的に求める方法

データストリームに対して、要素が次々と追加されていく状況を考えてみましょう。このようなストリームから、その時点までに追加されたすべての数値の中央値(メディアン)を常に効率よく取得できるシステムを実装する必要があります。

中央値とは、ソート済みのリストの中央に位置する値のことです。リストの長さが奇数の場合は中央の要素がそのまま中央値となり、偶数の場合は中央の2つの要素の平均を取ります。

この問題を解くために、次の2つのメソッドを実装します。

• addNum(): ストリームに数値を1つ追加する

• findMedian(): それまでに追加されたすべての数値の中央値を返す

アルゴリズム(解法の手順)

この問題は、2つの優先度付きキュー(priority queue)を使うことで効率的に解けます。片方を最大ヒープ(left)、もう片方を最小ヒープ(right)として機能させ、データをほぼ半分ずつ振り分けて管理します。

• 優先度付きキュー left(最大ヒープ)と right(最小ヒープ)を定義する

• addNum メソッドを定義する(引数は追加する数値 num):

  • left が空、または num が left の先頭要素より小さい場合は、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 を返す

この実装では、addNum の計算量は O(log n)、findMedian の計算量は O(1) となります。大量のデータが流れ込むストリーム処理において非常に効率的なアプローチです。

実装例

以下の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] なので、中央値は (10 + 15) / 2 = 12.5

• さらに 25 と 30 を追加すると [10, 15, 25, 30] となるため、中央値は (15 + 25) / 2 = 20

• 最後に 40 を追加すると [10, 15, 25, 30, 40] となるため、中央値は中央の要素である 25