遅延評価セグメント木をソラで書きたいあなたに

最下段が N-1 から始まるのと、親と子の取得方法がわかってればセグ木は書けます (ホンマか？)
— 009_tsutaj@CODE FESTIVAL (@_TTJR_) 2017年3月29日

この記事は前記事: 「セグメント木をソラで書きたいあなたに」の続編です。セグ木をソラで実装するのはまだ厳しい・・・という方はまずはこちらの記事を読んでみてください。

tsutaj.hatenablog.com

この記事は「セグメント木はソラで書けるけど、遅延評価セグ木は書けない・・・」という人を対象にしています。

遅延評価って何？

まず遅延評価について軽く説明しておきます。

遅延評価というのは、値の伝播*1を遅らせるテクニックです。これは区間に対して更新を行うときに必要になる時が多いです。例えば、値の更新を普通のセグメント木でやろうとしたとき、範囲の長さを $L$ とするとこの範囲の長さだけクエリを呼ばなければならないため $O(L \log N)$ かかってしまいますが、遅延評価をすることでこれを $O(\log N)$ で実現できます。

どうしたらうまくいくのか、簡単に解説します。

配列 a = {1, 3, 2, 6, 5, 4, 7, 9} を普通のセグメント木(区間和) で表すと、以下のようになります。

f:id:tsutaj:20170330223205p:plain

区間和に対応したセグメント木なので、親は子の 2 値の合計を持っている、といった作りになりますね。

このセグメント木に対して、 $\left[1, 5\right]$ 内の要素それぞれに 3 加算することを考えます。すると、通常ではこんなにも多くのノードを更新しなければなりません。

f:id:tsutaj:20170330223607p:plain

これでは効率が悪いですね。そこで、遅延評価用の配列を別に用意し、1 つのノードにつき 2 つの値を持つようにします。説明のため、普通のセグメント木にもあるような、値を覚える配列を「値配列」と、遅延評価用の配列を「遅延配列」と呼ぶことにします(造語)。遅延配列では何を持つのかというと、 自ノードの区間に一様に加えられた値 を持つことになります。

遅延配列をもつセグメント木は、このようなイメージになります。

f:id:tsutaj:20170330223212p:plain

例えば $\left[1, 5\right]$ に加算クエリが与えられた時、この区間を $\left[1, 4\right]$ と $\left[5, 5\right]$ にわけて考えます。この 2 つの区間に対応するノードについて、遅延配列を更新 します。(なぜなら、この 2 つの区間について一様に加算されているからです。)

f:id:tsutaj:20170330223214p:plain

遅延配列に値が入っているノードについては、いつか遅延配列の情報を値配列の方に伝播させなければいけません。これはどのタイミングで行えば良いかというと、「クエリでそのノードにアクセスしたとき」に更新すれば良いです。

f:id:tsutaj:20170330223217p:plain

遅延配列に値が入っているとき、値配列に入っている値は本来の値ではありません。遅延配列の値を加味してあげることで、初めて本来の値になります。ですので、そこにアクセスした時はまず「遅延配列の中身を反映させてから」処理を行う必要があるのです。

まとめると、

遅延配列は、区間に一様に与えられた値を記憶している
遅延配列に中身が入っているノードにアクセスした時、自分と自分の子に対して値の伝播が起こる
親ノードは、値配列が確定した状態の子ノードを参照することで更新される

このような実装をすることで、必要なときだけ値の伝播が起こるため計算量がぐっと減ります。これで、遅延評価付きのセグメント木も書けるようになります！

実際に書いてみよう

それでは「区間加算・区間和クエリ」を扱えるような遅延評価付きのセグメント木を書いてみましょう。

初期化

初期化については、普通のセグメント木とほぼ同じです。

というのも、最初に用意するときは、遅延配列に何も入っているはずがないですよね。実装は普通のセグメント木に遅延配列の宣言が増えただけです。

実際に書くとこのようになります。

struct LazySegmentTree {
private:
    int n;
    vector<ll> node, lazy;

public:
    LazySegmentTree(vector<ll> v) {
        int sz = (int)v.size();
        n = 1; while(n < sz) n *= 2;
        node.resize(2*n-1);
        lazy.resize(2*n-1, 0);

        for(int i=0; i<sz; i++) node[i+n-1] = v[i];
        for(int i=n-2; i>=0; i--) node[i] = node[i*2+1] + node[i*2+2];
    }
};

遅延評価

ノードに対して評価を行う関数を作りましょう。そのためには

自分のノードは何番目なのか？
自分のノードが指す範囲はどうなっているのか？

を知る必要があるため、 3 つの引数を取ります。

遅延配列に中身が入っている場合、評価を行う必要があります。この関数で何をすべきなのかというと、

自ノードの値配列に値を伝播させる
子ノードの遅延配列に値を伝播させる
自分のノードの遅延配列を空にする

これができればよいです。最下段のノードの処理に気をつけながら書くと、次のようになります。

// k 番目のノードについて遅延評価を行う
void eval(int k, int l, int r) {

    // 遅延配列が空でない場合、自ノード及び子ノードへの
    // 値の伝播が起こる
    if(lazy[k] != 0) {
        node[k] += lazy[k];

        // 最下段かどうかのチェックをしよう
        // 子ノードは親ノードの 1/2 の範囲であるため、
        // 伝播させるときは半分にする
        if(r - l > 1) {
            lazy[2*k+1] += lazy[k] / 2;
            lazy[2*k+2] += lazy[k] / 2;
        }

        // 伝播が終わったので、自ノードの遅延配列を空にする
        lazy[k] = 0;
    }
}

区間加算

区間に対しての更新なので、前記事でいうところの 最上段から下がっていく イメージで見ていきます。

基本的には前回の値の取得クエリとやることは同じです。しかし、遅延評価の関数を適切に呼び出す必要があります。

要求区間が対象区間を完全に被覆している場合、区間に対して一様に値が与えられていると言えます。ですので遅延配列に値を入れた後、評価関数を呼び出します。

そうでない場合、子ノードの値配列をしっかり更新してから、子ノードの値をもらって更新することをします。これは再帰関数にすることで簡潔に書くことができます。

実際に書くとこのようになります。

void add(int a, int b, ll x, int k=0, int l=0, int r=-1) {
    if(r < 0) r = n;

    // k 番目のノードに対して遅延評価を行う
    eval(k, l, r);

    // 範囲外なら何もしない
    if(b <= l || r <= a) return;
    
    // 完全に被覆しているならば、遅延配列に値を入れた後に評価
    if(a <= l && r <= b) {
        lazy[k] += (r - l) * x;
        eval(k, l, r);
    }

    // そうでないならば、子ノードの値を再帰的に計算して、
    // 計算済みの値をもらってくる
    else {
        add(a, b, x, 2*k+1, l, (l+r)/2);
        add(a, b, x, 2*k+2, (l+r)/2, r);
        node[k] = node[2*k+1] + node[2*k+2];
    }
}

区間和の取得

これも前回の区間和クエリと書き方はほぼ同じです。違いはやはり遅延評価を呼ぶかどうかで、今回の場合は関数の最初に遅延評価を入れるだけでよいです。

実際に書くとこのような感じになります。

ll getsum(int a, int b, int k=0, int l=0, int r=-1) {
    if(r < 0) r = n;
    if(b <= l || r <= a) return 0;

    // 関数が呼び出されたら評価！
    eval(k, l, r);
    if(a <= l && r <= b) return node[k];
    ll vl = getsum(a, b, 2*k+1, l, (l+r)/2);
    ll vr = getsum(a, b, 2*k+2, (l+r)/2, r);
    return vl + vr;
}

コード例

遅延評価セグメント木を習得することで様々なクエリを処理できるようになります。今回は AOJ の DSL で登場するクエリ処理のコード例を載せておきます。

※「区間更新」を扱うときは、遅延配列のフラグが立っているかどうかの bool 配列を別に用意しなければならないため、注意が必要です。

区間最大クエリにしたりなどは比較的簡単にできます。また、工夫すれば区間最小値の個数を答えるクエリなども作れますので、ぜひやってみてください。

楽しいセグ木ライフをお送りください。それではー。

*1:2018/05/19 追記: この記事で言う「伝播」とは、現在のノードに入っている値を子ノードに配ることを指します。子ノードに配られた値はさらにその子ノードに配られ、やがて最下段にあるノードまで順に配られるため、「伝播」という表現を用いています。

hogecoder

tsutaj 競技プログラミングの記録