日記 (2018 年 7 月 7 日)

セグメント木について勉強したのでまとめておきます。

長さ N の数列が与えられて、 以下の 2 操作を行う必要があるとします。

普通に数列を配列に保持してこの操作を行おうとすると、 数の置き換えは O(1) でできますが、 最小値を求めるのには最悪で O(N) かかります。 しかし、 セグメント木というデータ構造を用いることで、 両方の操作をそれぞれ O(logN) で行うことができます。 このアイデアとしては以下の通りです。

以下、 簡単のため N=2n と書けるとします。 まず、 0 番目と 1 番目そして 2 番目と 3 番目のように、 前から順に 2 つずつ対にして、 その 2 つの数の最小値を全て計算します。 そうすると、 2n1 個の数が計算されます。 次に、 この 2n1 個の数に対して、 同じように前から 2 つずつ対にして最小値を求めます。 これを順次繰り返していくと、 計算される数の個数がどんどん半分になっていき、 やがて 1 個になります。 最後に求まったこの数は数列全体の最小値となります。

図で表すと以下のような完全二分木になります。 これがセグメント木です。 四角形の左上に書いてある小さい数字は後で説明しますので、 とりあえず無視してください。

この木の特徴として、 あるノードに書かれている数は、 そこから下向きに枝を辿って到達できる葉ノード (つまり最初に与えられた数列の数が書かれたノード) 全体の最小値になっています。 このことを利用すると、 最初の数列の特定の範囲の最小値を高速に計算できます。

具体例として、 最初に与えられた数列の区間 [0,6) の最小値を求めてみます。 この区間の最小値は、 [0,4) の最小値と [4,6) の最小値の最小値であって、 [0,4) の最小値と [4,6) の最小値は下の図の緑に着色されたノードにすでに書かれています。 したがって、 この緑の 2 つのノード最小値を求めれば、 求めたかった [0,6) の最小値が求まります。

このように、 与えられた区間の最小値は、 セグメント木からうまくノードを見つけてくることで、 そのノードに書かれた数を使って高速に求めることができます。 問題はどのノードの数を使えば良いのかということですが、 これは以下のようにして求められます。

まず、 一番上のノードが使えるか判定します。 一番上のノードには [0,8) の最小値が書かれていますが、 この値は、 求めたい区間である [0,6) に入っていない余分な数とも最小値をとってしまっているので、 今回は不適です。 そこで、 次に 1 つ下のノードを見ます。 1 つ下の左側のノードは [0,4) の最小値が書かれていて、 これは求めたい区間の [0,6) に含まれているので、 この値は使えます。 したがって、 左側のノードの探索はこれで終了します。 一方で、 右側のノードは [4,8) の最小値が書かれていますが、 これは [0,6) からはみ出ているので不適切です。 したがって、 さらに枝を下ってもう 1 つ下のノードを見ます。 左側は [4,6) の最小値で、 これは [0,6) に含まれているので計算に用いることができて、 左側の探索を終了します。 一方で、 右側は [6,8) の最小値ですが、 これはもはや [0,6) と交わってすらいないので、 これ以上下方向に探索を続けても使える値はありません。 つまり、 右側の探索も終了して良いことになります。 以上で、 ノードの探索が終わったので、 使えると分かったノードの値全ての最小値を計算すれば、 求めたいものが得られます。

少し考えると、 この計算が O(logN) でできることが分かります。

次に、 最初に与えられた数列のどこかを別の数に書き換えることを考えます。 例えば、 2 番目の数を書き換えたいとします。 すると、 セグメント木の上の方にある数も書き換える必要が生じます。 しかし、 全てを書き換える必要はなく、 今回の場合は下の図の緑のノードを書き換えるだけで十分です。

書き換える箇所はセグメント木の高さの分だけなので、 この書き換え処理も O(logN) でできます。

さて、 N=2n と書けない場合はどうするかですが、 その場合はとりあえず N2n となる (最小の) n をとって、 数列の長さが 2n だと思って上のセグメント木を構築します。 余った箇所は、 最小値を取る操作に影響を与えない、 十分大きな数で埋めておけば問題ありません。

ところで、 セグメント木は別にある区間の最小値を求めたい場合でなくても使うことができます。 例えば、 ある区間の最大値を求めたい場合でも、 ある区間の総和を求めたい場合でも、 上の説明の 「最小値を求める」 という部分を考えたい操作に置き換えれば実現できます。 しかし、 考える演算が順序に依存する場合、 すなわち左から計算した場合と右から計算した場合で結果が異なり得る場合、 この構造は使えません。 つまり、 考える演算は結合的である必要があります。 また、 最初の数列の長さが 2 の冪でない場合は余った部分を何らかの数で埋める必要がありますが、 この埋める数は演算に影響を与えてはいけません。 言い換えれば、 考えている演算の単位元で埋める必要があります。 以上をまとめると、 結合的かつ単位的な演算が定まってさえいれば良いわけですが、 数学ではそういうものをモノイドと呼びます。 ということで、 セグメント木は、 値の更新とモノイド演算を対数時間で行うことができるデータ構造なわけです。

あとは実装をどうするかですが、 ほとんど上で説明したのをそのままコードに落とし込めば良いだけです。 ただし少し問題になるのが、 木をどう保持するかです。 セグメント木は完全二分木 (全ての葉ノードの深さが同じ) なので、 1 つの配列で保持することができます。 どうするかというと、 木の一番上のノードの値を最初の要素として格納し、 上から 2 番目の段のノードを左から順に続いて格納し、 次は上から 3 番目を、 とすれば良いです。 こうすると、 全体として長さ 2n+11 の配列に全ノードの値を保持できます。 この保持の方法だと、 インデックスが i のノードについて、 その親のノードのインデックスは (i1)/2 となり、 子のノードのインデックスは 2i+12i+2 となるので、 わりと簡単に親や子を参照できます。

なお、 上の図で、 各ノードの左上に小さく書かれていた番号は、 このように木を配列に保持した場合のインデックスを表していたのでした。

ということで、 Java での実装は以下のようになります。 ここでは、 与えられる数列は long で管理することにしました。 コンストラクタの unit には考えているモノイドの単位元を渡し、 operator にはモノイドの演算を渡します。

public static class SegmentTree {
private int size = 1;
private long[] array;
private long unit;
private Operator operator;
public SegmentTree(int n, long unit, Operator operator) {
this.unit = unit;
this.operator = operator;
while (size < n) {
size *= 2;
}
array = new long[2 * size - 1];
Arrays.fill(array, unit);
}
// i 番目の要素を取得する
public long get(int i) {
return array[i + size - 1];
}
// i 番目の要素を data に更新する
public void set(int i, long data) {
i += size - 1;
array[i] = data;
while (i > 0) {
i = (i - 1) / 2;
array[i] = operator.operate(array[i * 2 + 1], array[i * 2 + 2]);
}
}
// 区間 [a, b) に演算を施した結果を取得する
public long find(int a, int b) {
return find(a, b, 0, 0, size);
}
// セグメント木を格納した配列のインデックス i のデータが使えるかを調べる
// インデックス i のノードには区間 [l, r) に演算を施した結果が格納されている
public long find(int a, int b, int i, int l, int r) {
if (a >= r || b <= l) {
// 欲しい区間 [a, b) が調べている区間 [l, r) の完全に外にある場合
// ここの値は使わないので単位元を返しておく
return unit;
} else if (a <= l && b >= r) {
// 欲しい区間 [a, b) が調べている区間 [l, r) の完全に中にある場合
// ここの値は使えるのでそのまま返す
return array[i];
} else {
// 欲しい区間 [a, b) と調べている区間 [l, r) が微妙に重なっている場合
// 1 つ下のノードを見る
int m = (l + r) / 2;
// 子のノードは 2i + 1 と 2i + 2
return operator.operate(find(a, b, i * 2 + 1, l, m), find(a, b, i * 2 + 2, m, r));
}
}
}
@FunctionalInterface
public static interface Operator {
public long operate(long left, long right);
}

使うときはこんな感じにします。

int n = 100; // 数列の長さ
SegmentTree tree;
tree = new SegmentTree(n, Long.MAX_VALUE, Math::min); // 範囲内の最小値を取得
tree = new SegmentTree(n, Long.MIN_VALUE, Math::max); // 範囲内の最大値を取得
tree = new SegmentTree(n, 0, (a1, a2) -> a1 + a2); // 範囲内の総和を取得
tree.set(0, 3); // 値のセット
tree.find(1, 7); // 範囲内の演算結果を取得

さて、 ここまでの話は実は前座で、 本題はここからです。 上のセグメント木では、 範囲の演算結果の取得と 1 点の値の更新が O(logN) でできるわけですが、 値の更新もある範囲全てに一気に行いたいときもあります。 つまり、 以下の 2 操作を高速で行いたいわけです。

上の構造のままでやろうとすると、 更新は範囲の長さ分だけ行う必要があって、 結局最悪で O(NlogN) かかってしまい速さが足りません。 実は、 セグメント木を 2 つ作っておくと、 この場合でも両方の操作を O(logN) で行えます。 これについては、 また今度解説記事を作ろうと思います。