セグメントツリーの抽象化（C++）

※注終盤に「ARC008-D タコヤキオイシクナール」の解法のネタバレが存在します

■「セグメントツリーの抽象化ってしてる？」

●「抽象化……　セクハラですか、先輩」

■「ええっ」

●「冗談はさておいて……　抽象化というのは、セグ木自身に “区間和” や “指定要素を上書きして更新” などの具体的な情報を書かずに、それらを外部から与えることで、目的が変わっても同じコードを使えるようにすることですよね。はい、ある程度はしてますよ」

■「そうなんだ。いまいちやり方が分からないから、参考にしたいなと思って」

●「いいですよ、こんな感じです。”long long 型の値を持って区間minと一点上書きのクエリが与えられる” という場合、 “long long”、”区間min”、”一点上書き” という情報を外部から与えます」

template <class T>
class SegTree {
    int n;                       // 葉の数
    vector<T> data;              // データを格納するvector
    T def;                       // 初期値かつ単位元
    function<T(T, T)> operation; // 区間クエリで使う処理
    function<T(T, T)> update;    // 点更新で使う処理

    // 区間[a,b)の総和。ノードk=[l,r)に着目している。
    T _query(int a, int b, int k, int l, int r) {
        if (r <= a || b <= l) return def; // 交差しない
        if (a <= l && r <= b)
            return data[k]; // a,l,r,bの順で完全に含まれる
        else {
            T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2); // 左の子
            T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r); // 右の子
            return operation(c1, c2);
        }
    }

  public:
    // _n:必要サイズ, _def:初期値かつ単位元, _operation:クエリ関数,
    // _update:更新関数
    SegTree(size_t _n, T _def, function<T(T, T)> _operation,
            function<T(T, T)> _update)
        : def(_def), operation(_operation), update(_update) {
        n = 1;
        while (n < _n) {
            n *= 2;
        }
        data = vector<T>(2 * n - 1, def);
    }

    // 場所i(0-indexed)の値をxで更新
    void change(int i, T x) {
        i += n - 1;
        data[i] = update(data[i], x);
        while (i > 0) {
            i = (i - 1) / 2;
            data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);
        }
    }

    // [a, b)の区間クエリを実行
    T query(int a, int b) {
        return _query(a, b, 0, 0, n);
    }

    // 添字でアクセス
    T operator[](int i) {
        return data[i + n - 1];
    }
};

template <class T>

class SegTree {

int n; // 葉の数

vector<T> data; // データを格納するvector

T def; // 初期値かつ単位元

function<T(T, T)> operation; // 区間クエリで使う処理

function<T(T, T)> update; // 点更新で使う処理

// 区間[a,b)の総和。ノードk=[l,r)に着目している。

T _query(int a, int b, int k, int l, int r) {

if (r <= a || b <= l) return def; // 交差しない

if (a <= l && r <= b)

return data[k]; // a,l,r,bの順で完全に含まれる

else {

T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2); // 左の子

T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r); // 右の子

return operation(c1, c2);

}

public:

// _n:必要サイズ, _def:初期値かつ単位元, _operation:クエリ関数,

// _update:更新関数

SegTree(size_t _n, T _def, function<T(T, T)> _operation,

function<T(T, T)> _update)

: def(_def), operation(_operation), update(_update) {

n = 1;

while (n < _n) {

n *= 2;

}

data = vector<T>(2 * n - 1, def);

}

// 場所i(0-indexed)の値をxで更新

void change(int i, T x) {

i += n - 1;

data[i] = update(data[i], x);

while (i > 0) {

i = (i - 1) / 2;

data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);

}

// [a, b)の区間クエリを実行

T query(int a, int b) {

return _query(a, b, 0, 0, n);

}

// 添字でアクセス

T operator[](int i) {

return data[i + n - 1];

}

};

■「うーん、ちょっとややこしいね」

●「じゃあ1つずつ説明しますね」

テンプレートとstd::function

●「フィールド、すなわちメンバ変数はこんな感じですね」

template <class T>
class SegTree {
    int n;                       // 葉の数
    vector<T> data;              // データを格納するvector
    T def;                       // 初期値かつ単位元
    function<T(T, T)> operation; // 区間クエリで使う処理
    function<T(T, T)> update;    // 点更新で使う処理

    // 略
}

template <class T>

class SegTree {

int n; // 葉の数

vector<T> data; // データを格納するvector

T def; // 初期値かつ単位元

function<T(T, T)> operation; // 区間クエリで使う処理

function<T(T, T)> update; // 点更新で使う処理

// 略

}

■「template <class T> というのは、“T” を任意の型やクラスで置き換えられるということだね」

●「はい。今回は T は格納する値の型を表します。int, long long, double の他にも、行列などの自作クラスも入れられます」

■「行列を入れたかったら、class matrix を用意して、SegTree<matrix> st; みたいに宣言すれば良いんだね。この function<T(T, T)> っていうのは？」

●「これは、関数を変数として扱う std::function です。std::function<double(int, int)> と書くと、それは ” int 型の引数を2つ取って double 型の値を返す関数” を扱います。今は std を using しているので、std:: は省略しています」

■「へぇ、関数をこうやって扱えるんだね。今回の場合は、” T 型の変数2つを引数として、T 型の値を返す関数” になるのか。でもいまいちピンとこないな」

●「あとでまた出てくるので、そのときに詳しくお話しますよ。ちなみにこれを使うには、<functional> を include しないといけません」

■「def というのは、単位元だから、区間minクエリの場合は +∞、区間和クエリの場合は 0 ということだね」

●「はい。この実装では初期値も def で埋めています。初期値が単位元ではない場合は、初期化したあとに点更新クエリで全部書き換えないといけないですね」

■「簡単な場合だけに対応しているということだね」

コンストラクタ

●「コンストラクタはこうなっています」

    // _n:必要サイズ, _def:初期値かつ単位元, _operation:クエリ関数,
    // _update:更新関数
    SegTree(size_t _n, T _def, function<T(T, T)> _operation,
            function<T(T, T)> _update)
        : def(_def), operation(_operation), update(_update) {
        n = 1;
        while (n < _n) {
            n *= 2;
        }
        data = vector<T>(2 * n - 1, def);
    }

// _n:必要サイズ, _def:初期値かつ単位元, _operation:クエリ関数,

// _update:更新関数

SegTree(size_t _n, T _def, function<T(T, T)> _operation,

function<T(T, T)> _update)

: def(_def), operation(_operation), update(_update) {

n = 1;

while (n < _n) {

n *= 2;

}

data = vector<T>(2 * n - 1, def);

}

■「引数には、サイズの _n、初期値かつ単位元の _def に加えて、クエリ関数 _operation と更新関数 _update があるね」

●「_n はそのまま必要な要素数ですね。葉の数になるんですが、これが2冪でない場合は、_n以上で最小の2冪の数を葉の数として、メンバ変数 n にしています」

■「そして、葉の数の 2倍-1 の長さで vector<T> を用意して、初期値かつ単位元の def で埋めているんだね」

●「コンストラクタの引数リストのあとの : def(_def), operation(_operation), update(_update) は、コンストラクタ初期化子と呼ばれるものです。カッコの前のメンバ変数を、カッコの中身の値で初期化するということですね」

■「ここで初期値と必要な関数を渡しているんだね」

区間クエリ

●「肝心の区間クエリはこんな実装になっています」

template <class T>
class SegTree {
    // 区間[a,b)の総和。ノードk=[l,r)に着目している。
    T _query(int a, int b, int k, int l, int r) {
        if (r <= a || b <= l) return def; // 交差しない
        if (a <= l && r <= b)
            return data[k]; // a,l,r,bの順で完全に含まれる
        else {
            T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2); // 左の子
            T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r); // 右の子
            return operation(c1, c2);
        }
    }

  public:
    // [a, b)の区間クエリを実行
    T query(int a, int b) {
        return _query(a, b, 0, 0, n);
    }
}

template <class T>

class SegTree {

// 区間[a,b)の総和。ノードk=[l,r)に着目している。

T _query(int a, int b, int k, int l, int r) {

if (r <= a || b <= l) return def; // 交差しない

if (a <= l && r <= b)

return data[k]; // a,l,r,bの順で完全に含まれる

else {

T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2); // 左の子

T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r); // 右の子

return operation(c1, c2);

}

public:

// [a, b)の区間クエリを実行

T query(int a, int b) {

return _query(a, b, 0, 0, n);

}

■「_query() が本体だね。a, b がもとのクエリの範囲で、k が今注目しているノード番号、l, r がノード k の管轄範囲だね」

●「はい。l, r は k から分かるんですが、引数に残す方が分かりやすいかなと思ったので残しています」

■「無関係な範囲では単位元を返して、ノード k が区間 [a, b) に完全に含まれる場合は、そのノードの値を返すだけだね。大事なのは、範囲が交差する場合だね」

●「範囲が交差する場合は、まず左右の子に依頼して、それぞれが返す値 c1, c2 を得ます。この c1, c2 はもちろん T 型です」

■「区間minクエリ（RMQ）なら、min(c1, c2) を返して、区間和クエリ（RSQ）なら、c1+c2 を返せばいいんだね。そこで、あの function<T(T, T)> operation が出てくるのか」

●「まさにそうです！　区間ｍinクエリなら、 int func(int a, int b){ return min(a, b); } のような関数を operation に渡しておきます。区間和なら、int func(int a, int b){ return a + b; } を operation に渡しておきます」

■「どうやって渡すの？」

●「関数を作って渡す方法と、ラムダ式を使って渡す方法があります。他にもあるかもしれないです」

■「関数を作るのは、main関数の外に関数を作って、その関数名を渡すということだよね」

●「はい。それでもいいんですけど、関数の外に置かないといけないとか、関数名を考えないといけないとか、面倒な点があります。なので、1行で済む場合はラムダ式を使うことが多いですね」

■「ラムダ式って、聞いたことはあるけどどういうものなのかよく分かってないんだよね」

●「そんなに難しくないですよ。例えば int の和を返す関数なら、こう書きます」

[](int a, int b){ return a + b; }

■「普通の関数とあまり変わらないね。() の中身は引数、{} の中身は関数の処理みたいだけど、最初の [] は何？」

●「ここにはキャプチャと呼ばれる記述をします。このラムダ式の外にある変数をコピーして使う場合は [=]、参照して使う場合は [&] のように書きます。今回はキャプチャしない、つまりラムダ式の外にある変数を使わないので、詳しい話はググってください」

■「なるほど、この記述を引数としてそのまま入れたら良いの？」

●「はい。もちろん、こんなふうに変数に代入することもできます」

auto f = [](int a, int b){ return a + b; };

■「この場合、f の型は function<int(int, int)> になってくれるの？」

●「関数の戻り値の型はラムダ式に書いていませんが、return 文から推測して int と判断してくれて、function<int(int, int)> になってくれます。戻り値の型を明示する場合は、こう書きます」

[](int a, int b) -> int { return a + b; }

■「これで区間クエリは大丈夫だね。次は更新クエリか」

更新クエリ

●「更新クエリはこう実装しています」

    // 場所i(0-indexed)の値をxで更新
    void change(int i, T x) {
        i += n - 1;
        data[i] = update(data[i], x);
        while (i > 0) {
            i = (i - 1) / 2;
            data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);
        }
    }

// 場所i(0-indexed)の値をxで更新

void change(int i, T x) {

i += n - 1;

data[i] = update(data[i], x);

while (i > 0) {

i = (i - 1) / 2;

data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);

}

■「i += n – 1; の部分は、葉の添字を計算するところだね」

●「今回は、葉が n 個、葉以外が n-1 個あるので、n-1 を足せば i 番目の葉の添字が得られます」

■「そして、data[i] を更新するんだけど、update(data[i], x) という処理が呼ばれているね。これは？」

●「これは、複数の更新クエリに対応するための処理です。” i 番目に x を足す” と “i 番目を x に変える” の両方に対応しています」

■「なるほど、つまり、” i 番目に x を足す” なら、 update(data[i], x) は data[i]+x を返して、“i 番目を x に変える” なら、update(data[i], x) は data[i] を無視して x を返すんだね」

●「はい。update は、点加算なら [](T d, T x){ return d + x; }、点更新なら [](T d, T x){ return x; } になります。T の部分は実際に使う型に合わせます」

■「それで、葉を更新したら、そこから親をたどりながらさっきの operation で関連する葉以外のノードの値を更新していくんだね」

●「これでセグメントツリーの実装は完了です。私はついでに添字アクセスも用意してます」

添字アクセス

■「添字アクセスというのは、SegTree 型の変数 st に対して、st[i] で i 番目の値を得られるようにするということ？」

●「そうです。実装はこうなっています」

    // 添字でアクセス
    T operator[](int i) {
        return data[i + n - 1];
    }

// 添字でアクセス

T operator[](int i) {

return data[i + n - 1];

}

■「そのままだね。これだと、一点へのアクセスは O(1) になるね」

●「高速化はあまり意図してないんですけど、これがあるとデバッグが楽になります。出力ストリームへの書き出しも演算子オーバーロードすればもっと便利ですけど」

使用例

■「とりあえず、AOJのverify用問題をやってみようか」

●「はい。Range Query – Range Minimum Query (RMQ) を解いてみましょう。区間min、点更新で、最初はすべての値を 2^31-1 で初期化しています。2^31-1 がギリギリの値なので、型は long long にしちゃいましょう」

■「となると、引数はこんな感じかな。

_n : N

_def : (1LL<<31)-1

_operation : [](long long a, long long b) { return min(a, b); }（区間min）

_update : [](long long a, long long b) { return b; }（点更新）

●「それで大丈夫そうですね。main関数はこうなりますね」

int main(void) {
    int N, Q;
    cin >> N >> Q;
    // 要素数N・long long型・区間min・点更新
    SegTree<long long> st(N, (1LL << 31) - 1,
                          [](long long a, long long b) { return min(a, b); },
                          [](long long a, long long b) { return b; });
    rep(q, Q) {
        int c, x, y;
        cin >> c >> x >> y;
        if (c == 0) {
            // update(x, y)
            st.change(x, y);
        } else {
            // find(s, t)
            cout << st.query(x, y + 1) << endl;
        }
    }

    return 0;
}

int main(void) {

int N, Q;

cin >> N >> Q;

// 要素数N・long long型・区間min・点更新

SegTree<long long> st(N, (1LL << 31) - 1,

[](long long a, long long b) { return min(a, b); },

[](long long a, long long b) { return b; });

rep(q, Q) {

int c, x, y;

cin >> c >> x >> y;

if (c == 0) {

// update(x, y)

st.change(x, y);

} else {

// find(s, t)

cout << st.query(x, y + 1) << endl;

}

return 0;

}

■「よし、ACだね」

提出コード：http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=3451627#1

●「じゃあ、この勢いで ARC008-D タコヤキオイシクナールも解いてみましょう！」

タコヤキオイシクナール（解法ネタバレ注意）

■「この問題は、１次関数の合成について、結合法則が成り立って、単位元（1x+0）が存在するから、セグメントツリーに乗せられる、という問題だね」

●「はい。Nが非常に大きいため、座標圧縮が必要になりますが、今回はセグメントツリーの部分だけ確認しましょう」

■「タコヤキオイシクナールのボックスを表す struct box{ double a, b; } を用意して、2つのボックスを合成する関数をラムダ式で作って、こんな感じかな？」

struct box {
    double a, b;
};

int main(void){
    // 略（入力受け取りと座標圧縮）

    // 要素数N, 初期値box{1,0}, 区間合成, 点更新
    SegTree<box> st(N, box{1, 0},
                    [](box s, box t) {
                        return box{s.a * t.a, s.b * t.a + t.b};
                    },
                    [](box s, box t) { return t; });
    // ma: 最大値, mi: 最小値
    double ma = 1, mi = 1;
    rep(i, M) {
        st.change(p[i], box{a[i], b[i]}); // 点更新
        box whole = st.query(0, N);       // 全体の合成関数を計算
        double v = whole.a + whole.b;     // 全体に通したときの美味しさ
        ma = max(ma, v);                  // 最大値を更新
        mi = min(mi, v);                  // 最小値を更新　
    }
    // 略（出力）
}

struct box {

double a, b;

};

int main(void){

// 略（入力受け取りと座標圧縮）

// 要素数N, 初期値box{1,0}, 区間合成, 点更新

SegTree<box> st(N, box{1, 0},

[](box s, box t) {

return box{s.a * t.a, s.b * t.a + t.b};

[](box s, box t) { return t; });

// ma: 最大値, mi: 最小値

double ma = 1, mi = 1;

rep(i, M) {

st.change(p[i], box{a[i], b[i]}); // 点更新

box whole = st.query(0, N); // 全体の合成関数を計算

double v = whole.a + whole.b; // 全体に通したときの美味しさ

ma = max(ma, v); // 最大値を更新

mi = min(mi, v); // 最小値を更新　

}

// 略（出力）

}

●「ACですね！　完璧です！」

提出コード：https://atcoder.jp/contests/arc008/submissions/4784678

■「これ1つで色んなタイプのセグメントツリーを作れるのは便利だね」

●「はい！　渡す関数を間違えないようにしないといけないですけど、考えるところがそこだけなので、バグも生まれにくいですね」

■「遅延伝播セグメントツリーも抽象化してるの？」

●「一応してますけど、そっちはもっとややこしいので、また別の機会にお話しますね」

■「とりあえずいくつかのパターンで抽象化セグメントツリーを使ってみて慣れようかな」

Pocket

セグメントツリー入門より:

2019年6月10日 10:59 AM

[…] ■「AOJ には、セグメントツリーで区間和を計算する問題があるから、マスターになるならこっちも考えてみる必要があるよ。また、セグメントツリーを複数使うときのために class 化しておくべきだし、GCD みたいな、結合法則が成り立つけど逆元が無いような演算も使えるようになると良いかも。タコヤキオイシクナールは特殊なモノイドをセグ木に乗せる問題の代表だし、汎用性を高めるためにも抽象化をしておくと良いかも。Young Maids みたいな問題のために最小要素の添え字を返す形も作っておいた方が良いね。それから遅延伝播セグメントツリーというのもあって……」 […]

返信