Lilliput Steps

小さな一歩から着実に. 数学やプログラミングのことを書きます.

Introduction to Segment Tree

こんにちは, kagamiz(@kagamiz) です!
この記事は, Competitive Programming Advent Calendar Div2012 18日目の記事として書かれました.
この記事では, ぼくが気に入っているデータ構造であるセグメント木について, 説明と問題演習・問題紹介を交えながら紹介していきたいと思います.

RMQ(Range Minimum Query)を処理するセグメント木はわかるけど, そこから応用するのが難しい...という方向けに書いていきます.

説明の章の先頭には, 水色の★もしくは黄色の★をつけています. は「プログラミングコンテストチャレンジブック」などでも触れられている事項で, はそうでない事項を指します.
演習問題の先頭には★マークをつけていますが, は基本的な問題, はやや難しい問題として示しています.

概要についてですが, ここではiwiwi先生が書いてくれた「プログラミングコンテストでのデータ構造」の33ページ目からの事項を参照してください.(ここに書いてあることを前提に説明を進めていきます.)

まずは, 上の記事を参考に, 下の3問を解いてみましょう.

演習1. Range Minimum Queryを解いてください.
演習2. Range Maximum Queryを解いてください.
演習3. Big Range Queryを解いてください.

和を持つセグメント木とBIT
Partial Sum Problem を解くセグメント木を考えてみましょう.

この問題で扱うセグメント木は, セグメント木のノードに, そのノードに対応する区間の和を持つと実現できそうです.
この問題を解くために, 「1 点に対して行う処理」と「大きい区間で受け止める処理」を RMQ のときと同様に考えてみましょう.

・1点に対して行う処理
この処理に当たるのは, addクエリ(場所xに値valを足すクエリ)です.
これは, RMQ の update 関数を少し変形するだけで実現できます.

void add(int x, int val)
{
    x += size - 1;
    
    seg[x] += val; //場所kに対応する節点にxを足す.
    
    while (x){
        x = (x - 1) / 2; //親に戻る
        seg[x] = seg[x * 2 + 1] + seg[x * 2 + 2]; //親の区間の和は, 子の区間の和を足しあわせたもの.
    }
}

・大きい区間で受け止める処理
この処理は, sumクエリ(区間[l, r]の和を求めるクエリ)に当たります.
これも, RMQ の getMin 関数をすこし変形するだけで実現できます.

//区間[a, b)の和を求める.
int getSum(int a, int b, int k, int l, int r)
{
    if (r <= a || b <= l){ //もし区間が交わっていなければ
        return (0); //0を返す
    }
    
    if (a <= l && r <= b){ //[l, r)が[a, b)に完全に内包されていれば
        return (seg[k]); //そのノードを返す
    }
    
    return (getSum(a, b, k, l, (l + r) / 2) + getSum(a, b, k, (l + r) / 2, r)); //区間を2つに分割して和を求める
}

これで, 和を持つセグメント木を構築することができます.

しかし, 和を持つセグメント木は, このように実装しなくても, BIT(Binary Indexed Tree; Fenwick Tree)というデータ構造を用いることで, 軽く実装することができます.
(BITもセグメント木の兄弟のようなデータ構造です.)

BITは,
・数列の要素A[1]からA[x]までの和を求める. (=sum(x))
・場所xに値valを加算することができる. (add(x, val))
データ構造です.

このクエリが実装できれば, 最初の問題での[l, r]の和を求める処理を, sum(r) - sum(l - 1)で求めることができます.
以下に, C++を用いた実装例を示します. (なぜこの様にして動くのかは, プログラミングコンテストチャレンジブックを参照してください.)

int bit[N + 1];

//BITは 1-indexed で扱います!!!! [1, N] 区間を配列の添字として扱うときに注意.
void add(int x, int val)
{
    while (x <= N){
        bit[x] += val;
        x += x & -x;
    }
}

int sum(int x)
{
    int ret = 0;
    while (x){
        ret += bit[x];
        x &= (x - 1);
    }
    
    return (ret);
}

演習4. Partial Sum Problemを解いてください.
演習5. Partial Sum Problemを BIT を使って解いてください.
演習6. RMQ に於いて, 最小(大)値クエリが[1, x]の形だけをとり, 更新後の値が更新前より必ず大きくなるのであれば, RMQ を BIT で実装することができます. そのときの RMQ を, BITで実装してみてください. (このアイディアはsnukeさんから教わりました.ありがとうございます!)

セグメント木の応用例(1)
セグメント木を応用した問題の中に, 次のようなものがあります(A Simple Problem with Integers).
この問題では, 次のようなクエリを処理しなければいけません:
区間[l, r]に数xを足す.
区間[p, q]の和を求める.

やはり区間を扱うのでセグメント木が適していそうです. しかし, 前回までのように, 対応する区間全てにxを加えると, 時間がかかってしまいます.

そこで, 次のように2つのセグメント木を持つことを考えます.
1. 区間[l, r]に一様に加算される値を持つセグメント木 - (allという配列と仮定しましょう
2. 区間[p, q]に部分的に加算される値を持つセグメント木 - (partという配列と仮定しましょう

こうすると, 区間[p, q]に対応する節点の番号がkであるとき, [p, q]の和は (q - p + 1) * all[k] + part[k]で求まります.

また, 区間[l, r]に値xを足す時, 区間[l', r']が完全に区間[l, r]に内包されていれば, all[k]にxをたし, そうでない場合はpart[k]に, 交わっている区間の大きさ * xを足してやれば良いです.

実装例を以下に示します.

static const int MAX_SIZE = 1 << 17; //segment tree のサイズ。 2^17 ≒ 1.3 * 10^5

typedef long long Int;
Int all[2 * MAX_SIZE - 1], part[2 * MAX_SIZE - 1]; // segment tree

//区間[a, b)に値xを加算する.
void add(int a, int b, int x, int k = 0, int l = 0, int r = MAX_SIZE)
{
    if (a <= l && r <= b){ //[l, r)が[a, b)に完全に内包されていれば
        all[k] += x; //[l, r)の全ての区間が持つ値としてxを足す.
    }
    else if (l < b && a < r){ //[l, r)と[a, b)が交差していれば
        part[k] += (min(b, r) - max(a, l)) * x;  //交差している分の値を, 部分的な和を持つノードに加算する.
        add(a, b, x, k * 2 + 1, l, (l + r) / 2); //子でも同じ処理を行う.
        add(a, b, x, k * 2 + 2, (l + r) / 2, r); //〃.
    }
}

Int sum(int a, int b, int k = 0, int l = 0, int r = MAX_SIZE)
{
    if (b <= l || r <= a){ //[a, b)と[l, r)が全く交差しない場合
        return (0);
    }
    else if (a <= l && r <= b){ //完全に内包されていれば
        return (all[k] * (r - l) + part[k]);
    }
    else { //[l, r)と[a, b)が交差していれば
        ll res;
        res = (min(b, r) - max(a, l)) * all[k]; //そのノードの全ての要素が持つ値のうち, [a, b)に属すものの分だけを加算する.
        res += sum(a, b, k * 2 + 1, l, (l + r) / 2); //子ノードで和を求める.
        res += sum(a, b, k * 2 + 2, (l + r) / 2, r); //〃
        return (res);
    }
}

なお、この問題はBITを使って解くことも出来ます. ゆっくり考えてみましょう.

演習7. A Simple Problem with Integersを解いてください.
演習8. A Simple Problem with IntegersをBITを使って解いてください.

セグメント木の応用例(2) (Starry Sky木)

さて, 先程示したように, 役割を2つのセグメント木で分離する, という考え方を用いると, 次のようなクエリを処理できるセグメント木を構築することも可能です:

区間[l, r]に値 val を加算する.
区間[l, r]の最大(小)値を求める.

やはり区間を扱うので, セグメント木が適していますね. この処理をする問題(Range Minimum Query 2)を解いてみましょう.

ここでは、
・その区間に一様に加算される値を持つセグメント木(segAddという配列とします
・子の区間の最大(小)値を持つセグメント木(segMinという配列とします.

をもつと, ある区間[l, r]の最小値は, その区間に対応する節点の番号をkとすると, segMin[k] + segAdd[k]となります.

これも, 実装例を以下に示します.

static const int MAX_SIZE = 1 << 17; //segment tree のサイズ。この実装では2べきにする必要がある。 2^17 ≒ 1.3 * 10^5

typedef long long Int;
Int segMin[2 * MAX_SIZE - 1], segAdd[2 * MAX_SIZE - 1];

//区間[a, b)に値xを加算する.
void add(int a, int b, int x, int k = 0, int l = 0, int r = MAX_SIZE)
{
	if (r <= a || b <= l) return; //もし交差しない区間であれば終える.
	
	if (a <= l && r <= b){ //もし今みている区間[l, r)が[a, b)に完全に内包されていれば
		segAdd[k] += x;  //区間[l, r)にkを加算する.
		return;
	}
	
	add(a, b, x, k * 2 + 1, l, (l + r) / 2); //子の区間に(必要があれば)xを加算する.
	add(a, b, x, k * 2 + 2, (l + r) / 2, r); //〃

	//親の区間の最小値は, 子の区間の最小値 + 自分に一様に加算されている値 である.一様に加算される値は更新しなくて良い.
	segMin[k] = min(segMin[k * 2 + 1] + segAdd[k * 2 + 1], segMin[k * 2 + 2] + segAdd[k * 2 + 2]);
}

Int getMin(int a, int b, int k = 0, int l = 0, int r = MAX_SIZE)
{
	if (r <= a || b <= l) return (LLONG_MAX);
	
	if (a <= l && r <= b) return (segMin[k] + segAdd[k]); //完全に内包されていれば,その区間の最小値を返す.
	
	lint left = getMin(a, b, k * 2 + 1, l, (l + r) / 2); //子の区間の最小値を求める.
	lint right = getMin(a, b, k * 2 + 2, (l + r) / 2, r); //子の区間の最小値を求める
	
	return (min(left, right) + segAdd[k]); //親の区間の最小値は, 子の区間の最小値 + 自分に一様に加算されている値 である (大切なので2回書きました!!)
	
}


この様な処理ができるセグメント木には名前がついており, Starry Sky木と呼ばれています(JOI春合宿のStarry Skyという問題が起源とされているため.)
この他にも色々な実装方法があります.

また, セグメント木の他のテクニックとして, 遅延評価があります. こちらはきゅうりさんの記事で詳しく解説されています.
今述べた応用例(1)(2)は, それぞれセグメント木の遅延評価を用いても実装できます. どちらの手法でもかけますが, 遅延評価のほうも応用範囲が広いので, 是非書いてみましょう!

演習9. Range Minimum Query 2 を解いてください.

問題紹介
セグメント木やBITを使った問題は, プログラミングコンテストチャレンジブックや, とざんさんの記事に書かれています.
ここの問題を解いて, セグメント木をマスターしていきましょう!

ここに載っている問題の他に, 以下の問題をお勧めします. いずれの問題も, この記事で触れたセグメント木を応用して解くことが可能です. セグメント木, やっぱりすごいですよネ!

Dungeon(AOJ 0553) (高速なO(N)解法アリ)
School of Killifish(AOJ 1068)
The Incubator(AOJ 1083)
House Moving(AOJ 2431)
Bookshelf(2011年 JOI春合宿4日目問2)
Circular RMQ(Codeforces 52C)
Buses(Codeforces 101B)
だるま落とし(第2回早稲田大学プログラミングコンテスト)