はじめに

セグメント木と遅延評価セグメント木を自分がどうやって空で実装しているのかをメモする。C言語での実装。セグメント木が実装できれば、大体の場面で平衡二分探索木を実装する必要性がなくなり、競技プログラミング等で解ける問題の幅がぐっと広がる。

また、Binary Indexed Tree(BIT)の上位互換なので、とりあえずセグメント木さえきちんと実装できれば、BITとかをあまり気にしなくて良くなる。

Reference

http://tsutaj.hatenablog.com/entry/2017/03/30/224339
https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/1
蟻本のセグメント木の解説

AOJ courceの問題中心に例題もこの記事の一番下の章に挙げている。

セグメント木の実装のコツ

priority queueと似たようなデータ構造なので、まずはpriority queueを自力で実装できるかをチェックする。C++だとpriority_queueのライブラリがあるけど、残念ながらC言語だとないので自分は空で書いてる。priority queueの解説はプログラミングコンテスト攻略のためのアルゴリズムとデータ構造の10章に詳しい。

今回はpriority queueを紹介したいわけではないので、実装は飛ばす。自信がなければ例題や上の本の解説を見てほしい。

priority queueが書ければ、セグメント木の実装はそんなに難しくないし、自然に実装できる。以下はRAQ(range add query)の例：

struct range {
    int start;
    int end;
};

int is_overlap(struct range r1, struct range r2) {
    return r1.end > r2.start && r2.end > r1.start;
}

int is_contain(struct range this, struct range in) {
    return in.start <= this.start && this.end <= in.end;
}

void add(int *seg, struct range r, int val, int node, struct range this) {
    if(is_contain(this, r)) {
        // do something
        seg[node] += val;
    } else if (is_overlap(r, this)) {
        struct range left = {this.start, (this.start+this.end)/2};
        struct range right = {(this.start+this.end)/2, this.end};
        add(seg, r, val, node*2, left);
        add(seg, r, val, node*2+1, right);
    }
    // if not overlap, do nothing
    return;
}

int get(int *seg, int num, int idx) {
    int node = num + idx;
    // If necessary, do something on seg[node]
    int ans = 0;
    while(node) {
        // do something
        seg[node] += val;
        node /= 2;
    }
    return ans;
}

以下のように使う。

#define ROOT_NODE 1

int main(void) {
    int num, queries;
    int num2 = 1;
    scanf("%d %d", &num, &queries);
    while(num > num2) num2 *= 2; // power of 2

    static int seg[NUM2_MAX*2];
    struct range whole = {0, num2}; // [0, num2)
    // val: input // r: struct range; [start, end)
    add(seg, r, val, ROOT_NODE, whole);
    // or
    // idx: input
    int ans = get(seg, num2, idx);
}

個人的にはpriority queueの実装と同様にsegment treeを1-indexedにしたほうがセグメント木のノードの操作がしやすいので嬉しい。例えば、親ノードにアクセスするには、nodeの偶奇性を気にすることなく、node/2とすればよい。またセグメント木の配列の値は[num, 2*num)なので、実際の配列のindexにsizeを足す(つまり、seg[size+index]とする)だけでよい。

後、セグメント木は引数が多くなる傾向にあるので、struct rangeでまとめていて実装の見通しが良いと思う。

struct rangeを閉区間で表現するか、半開区間で表現するかは正直どちらでも良いが、半開区間で表現すると区間上の配列の個数をrange.end - range.startと表現でき、+1とか-1とかの余計な数がでないのが嬉しいので、自分は半開区間で実装する。

その他、注意すべき点

segの配列数は2乗で2*NUM2_MAX(入力配列数を2乗にして、その数の2倍)とする
Range Update Query(RUQ)などupdate系のやつを実装する場合は、初期化の際に架空の値を入れておく。これでその値がすでにupdateされたかどうかが判定できる。
is_overlapやis_containはそれぞれ、区間同士がoverlapしているか、区間同士が包含しているかを見る関数だが、間違って実装しやすいので注意する。(難しくはないけど、よくtypoしてしまう)

遅延評価セグメント木の実装のコツ

遅延評価セグメント木の書き方には蟻本の実装のようなやり方と、tsutajさんが提案している方法(遅延評価の部分を分離する方法)があるが、後者の方が応用が効き統一性のある実装ができるので、この記事では後者に絞って実装のコツを述べる。

考え方自体は http://tsutaj.hatenablog.com/entry/2017/03/30/224339 を見るのが一番良い。

こちらも、セグメント木の実装とにているが、RSQ and RAQを例に取る：

struct pair {
    int64_t sum;
    int64_t lazy_sum;
};

struct range {
    int start;
    int end;
};

#define NUM_MAX 100000
#define NUM2_MAX (1<<17)

#define ROOT_NODE 1

int is_overlap(struct range r1, struct range r2) {
    return r1.end > r2.start && r2.end > r1.start;
}

int is_contain(struct range elem, struct range box) {
    return elem.start >= box.start && elem.end <= box.end;
}

void unlazy(struct pair *seg, int node, struct range nr) {
    if(!seg[node].lazy_sum) return;
    int64_t lazy = seg[node].lazy_sum;
    seg[node].sum += lazy;
    seg[node].lazy_sum = 0;
    // if not simple
    if(nr.end - nr.start > 1) {
        seg[node*2].lazy_sum += lazy / 2;
        seg[node*2+1].lazy_sum += lazy / 2;
    }
}

void add_range(struct pair *seg, int64_t val, struct range r, int node, struct range nr) {

    // First of all, the lazy value should be woken
    unlazy(seg, node, nr);
    if(!is_overlap(r, nr)) return;
    if(is_contain(nr, r)) {
        seg[node].lazy_sum += val * (nr.end-nr.start);
        unlazy(seg, node, nr);
        return;
    }
    struct range left = { nr.start, (nr.start+nr.end)/2 };
    struct range right = { (nr.start+nr.end)/2, nr.end };
    add_range(seg, val, r, node*2, left);
    add_range(seg, val, r, node*2+1, right);
    // update
    seg[node].sum = seg[node*2].sum + seg[node*2+1].sum;
}

int64_t get_sum(struct pair *seg, struct range r, int node, struct range nr) {

    // First of all, the lazy value should be woken
    unlazy(seg, node, nr);
    if(!is_overlap(r, nr)) return 0;
    if(is_contain(nr, r)) return seg[node].sum;

    struct range left = { nr.start, (nr.start+nr.end)/2 };
    struct range right = { (nr.start+nr.end)/2, nr.end };
    int64_t s1 = get_sum(seg, r, node*2, left);
    int64_t s2 = get_sum(seg, r, node*2+1, right);
    return s1 + s2;
}

セグメント木と同じく、配列の個数を2*NUM2_MAX(入力配列数を2乗にして、その数の2倍)とし、クエリの区間は半開区間、つまり、[a, b)としている。

int segの代わりにstruct pairでlazy_sumが登場している。lazyの意味合いはその区間全てに加算される値という意味。unlazyという関数で子のノードに値を伝搬している。unlazyを実行した後は、sumの値が確定するので、ある区間の合計値はseg[node].sumで取得できる。

add_range関数は子ノードに降りていって、add_range(seg, val, r, node*2, left);とadd_range(seg, val, r, node*2+1, right);を実行し、sumの値を次々と確定させる。一番下まで行ったらseg[node].sum = seg[node*2].sum + seg[node*2+1].sum;として、親ノードの区間のsumの値を確定させる。

この実装の大きなポイントは紛れもなくunlazyという関数で、必要になった区間に関して遅延的に、lazyの値をsumの値に還元する。これはクエリ(get_sumやadd_range)ごとに必ず呼ばれる必要があるので、関数の頭に指定している。また、add_range関数では、lazy_sumが更新されれば、その都度unlazy関数によって、lazyの値をsumの値に還元するようにしている。

実装のポイントだが、自分の頭の中では、

普通にセグメント木を組む + unlazyでlazyの値をsumの値に遅延還元する(遅延評価のコアの部分)

という要領で実装している。この実装には2つの利点がある。

RMQ(Range Min Query)だろうがRUQ(Range Update Query)だろうが、unlazyという遅延評価する関数をカスタマイズすれば、後はセグメント木を実装するのとほとんど変わらない。
遅延評価する部分の関数と実際に評価する関数がほぼ分離されている。get_sumでは値を取得する際は、sumというメンバしか用いていない。これはunlazyという遅延評価する関数がsumの値を確定させてくれるからだ。蟻本の解答例だと分離されていない。

実際に遅延評価セグメント木の例題をやってみると、4つとも「セグメント木 + unlazy」のフレームワークに乗せて遅延評価セグメント木の実装を書けると思う。

以下、細かな補足

unlazy関数では、親のlazyの値を子に半分にして(lazy / 2)伝搬させているが、これはRange Add Query(RAQ)に依存するロジックであり、Range Update Query(RUQ)などのupdate系だと普通にlazyの値を更新するだけで良い。また、if(seg[node].laze == INIT_VAL) return;として、更新されていない場合は早期リターンするように初期値に架空の値INIT_VALを与えて区別する。詳しくは、DSL_2_I RSQ and RUQの解答例を見ると良い。

例題

priority queue

ダイクストラ法やクラスカル法でも計算量を落とすためにpriority queueが併用されるが、今回は純粋にpriority queueメインの問題を挙げている。

問題	解答例
ALDS1_9_C priority queue	最大ヒープ, ALDS1_9_Bは導入問題
ALDS1_15_D ハフマン符号	最小ヒープ, DFS
2017 CODE FESTIVAL THANKS C - Factory	最小ヒープ
ABC 062 D - 3N Numbers	最小ヒープ

セグメント木

問題	解答例
DSL_2_A Range Minimum Query(RMQ)	セグメント木
DSL_2_B Range Sum Query(RSQ)	セグメント木
DSL_2_D Range Update Query(RUQ)	更新順をセグメント木に持たせる
DSL_2_E Range Add Query(RAQ)	セグメント木

遅延評価セグメント木

問題	解答例
DSL_2_F RMQ and RUQ	遅延評価セグメント木
DSL_2_G RSQ and RAQ	遅延評価セグメント木
DSL_2_H RMQ and RAQ	遅延評価セグメント木
DSL_2_I RSQ and RUQ	遅延評価セグメント木