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

セグ木がソラで書けなかったらセグ木に何か生やす問題とか解けない
— つたじろう@帰省 (@_TTJR_) March 29, 2017

セグ木にいろいろ生やす問題がてんで解けない私なので、セグ木に慣れようと思い立ちました。そのためにはまずセグ木をもっと知らねばならないと思ったので、ソラで書けないか色々やっていました。

やってくうちにコツを掴んでソラで書けるようになってきたので、忘れないためにも記事を書くことにしました。ということで、何も見ずにセグ木を書くために押さえておきたいポイントを紹介していきます。

注意

この記事は、「セグ木がどんな形をしているかはわかるけど、実装はできない・・・」という方向けで、遅延評価が不要な、区間最小の基本的なセグメント木のみ扱っています。

そもそもセグ木って何という人は他記事を参考にしてください。
おすすめは iwiwi さんのスライド → プログラミングコンテストのためのデータ構造

また、すでにセグ木が書けるプロの方が見ても面白くないかもしれません。

記事では区間最小についてのみ扱っていますが、区間和もほぼ同じコードで実現できます (コード例は後述)。

まずは形を見よう

ノードの数

セグメント木は皆さんご存知のとおり、図で書くとこんな形をしています。

f:id:tsutaj:20170329204251p:plain

元のデータ数以上になる最小の 2 冪の数を $N$ とします。すると、ノードは全部で $2N - 1$ 個あります。なぜこうなるかは、各段にあるノードの数を 2 進で書いて足していくと簡単にわかります。

すなわち、

最下段にはノードが $N$ 個ある
それより上の段のノードは全部で $N-1$ 個ある

ということになります。この性質は後に大事になるので、しっかり把握しておきましょう。

ノードの関係

セグメント木上で操作をするには、ノードの関係を明らかにする必要があります。ここでは木と同様に、自分のノードの直下にあるノードを「子」、自分のノードの真上にあるノードを「親」と呼ぶことにします。

木に例えるとこれは完全二分木になります。つまり、葉でない限り自分の子は必ず $2$ 個あるということです。

f:id:tsutaj:20170329204255p:plain

自分のノードの番号を $k$ をおくと、親や子にアクセスするには以下のようにします。

親にアクセスするには、 $\lfloor \frac{k-1}{2} \rfloor$ 番目にアクセスする
子にアクセスするには、 $2k+1$ 番目・ $2k+2$ 番目にアクセスする

これを各所で使うことになります。

更新・取得クエリ

一般的なセグメント木では、値を更新するクエリと、値を取得するクエリの 2 種類を行うことになります。かなりざっくりとしていますが、各クエリに関してはこのようなイメージを持つと良いです。

値の更新クエリは、最下段から上がっていくようにする
値の取得クエリは、最上段から下がっていくようにする

雑ですが、このようなイメージを持って実装を重ねることで詰まることなく書けるようになると思います。

実際に書いてみよう

初期化

元となる配列があって、それをセグメント木で表してみることをやってみましょう。

まずは、セグメント木のサイズ (ノード数)を決定します。そのためには最下段のサイズを求める必要がありますが、これは先述のとおり、元のサイズ以上になる最小の 2 冪 です。この値を $N$ としたとき、セグメント木のサイズは $2N-1$ です。

次に実際に値を入れていくのですが、最下段から値を入れていき、それ以降は下の段から順番に自分の子を参照することで値を入れていきます。上のノードの値を決めるには下のノードを参照しなければならないことを考えると、最下段から入れるのは自然な流れですね。

最下段のノードのインデックスはどうなるのか？と思うかもしれませんが、これには最下段より上の段のノードは全部で $N-1$ 個ある ことを利用します。前に $N-1$ 個の要素があるので、インデックスは元のインデックスに $N-1$ を足せばよいですね。

これらをまとめると、次のような実装になります。

struct SegmentTree {
private:
    int n;
    vector<int> node;

public:
    // 元配列 v をセグメント木で表現する
    SegmentTree(vector<int> v) {
        // 最下段のノード数は元配列のサイズ以上になる最小の 2 冪 -> これを n とおく
        // セグメント木全体で必要なノード数は 2n-1 個である
        int sz = v.size();
        n = 1; while(n < sz) n *= 2;
        node.resize(2*n-1, INF);

        // 最下段に値を入れたあとに、下の段から順番に値を入れる
        // 値を入れるには、自分の子の 2 値を参照すれば良い
        for(int i=0; i<sz; i++) node[i+n-1] = v[i];
        for(int i=n-2; i>=0; i--) node[i] = min(node[2*i+1], node[2*i+2]);
    }
};

値の更新

「 $x$ 番目の要素を $val$ に更新する」ことを考えます。

これを行うには、 $x$ 番目の要素が含まれる区間全てを更新する必要があります。上のノードを更新するには下のノードを見なければならないため、下のノードから更新していかなければいけないことがわかります。

最下段から上がっていく イメージで書いていきます。つまり、まずは最下段を更新し、あとはその親を更新することを繰り返せばよいです。

実際に書くとこんな感じになります。

void update(int x, int val) {
    // 最下段のノードにアクセスする
    x += (n - 1);

    // 最下段のノードを更新したら、あとは親に上って更新していく
    node[x] = val;
    while(x > 0) {
        x = (x - 1) / 2;
        node[x] = min(node[2*x+1], node[2*x+2]);
    }
}

値の取得

「区間 $[a, b)$ にある要素の最小値を答える」ことを考えます。

説明のため、クエリで与えられた区間 (実際に計算したい区間) を「要求区間」、各ノードがカバーしている区間を「対象区間」と定義します (造語)。

これを扱うには、以下の 3 つの情報が必要になります。

要求区間はどのような区間か？
いま自分がいるノードは何番目か？
自分がいるノードはどのような区間か？ (対象区間はどのようになっているか？)

要求区間と対象区間の関係によって場合分けをします。

要求区間と対象区間が交わらない場合

この場合は、これ以上操作をすすめても意味がありませんので、答えに影響しない値を適当に返しておきます。

要求区間が対象区間を完全に被覆している場合

この場合、対象区間は要求区間の計算に必要です。そのため、現状の答えと比較して、更新が必要ならば更新をしていきます。

要求区間が対象区間の一部を被覆している場合

この場合、対象区間の子にあたる区間に移動して、完全に被覆するまで操作を行わなければなりません。

子に移動するので、「現在見ているノード」と「対象区間の情報」が変わります。子のノードのインデックスの取得は先述のとおりで、対象区間に関しては、子のノードが現在見ているノードを半分に分割したものであることを利用すると得られます。

最上段から下がっていくイメージでこれをまとめると、このようなコードになります。(ここが一番難しいと思いますので、わからなければ Twitter などで私に質問してください)

// 要求区間 [a, b) 中の要素の最小値を答える
// k := 自分がいるノードのインデックス
// 対象区間は [l, r) にあたる

int getmin(int a, int b, int k=0, int l=0, int r=-1) {
    // 最初に呼び出されたときの対象区間は [0, n)
    if(r < 0) r = n;

    // 要求区間と対象区間が交わらない -> 適当に返す
    if(r <= a || b <= l) return INF;

    // 要求区間が対象区間を完全に被覆 -> 対象区間を答えの計算に使う
    if(a <= l && r <= b) return node[k];

    // 要求区間が対象区間の一部を被覆 -> 子について探索を行う
    // 左側の子を vl ・ 右側の子を vr としている
    // 新しい対象区間は、現在の対象区間を半分に割ったもの
    int vl = getmin(a, b, 2*k+1, l, (l+r)/2);
    int vr = getmin(a, b, 2*k+2, (l+r)/2, r);
    return min(vl, vr);
}

コード例

AOJ にコードを上げていますので、参考までに。

以上です。遅延評価付きのセグ木についても後日記事でまとめます (たぶん)。

hogecoder

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