怖くない！　抽象化動的セグメント木の非再帰実装

はじめに

区間取得クエリを対数時間で処理するセグメント木は競技プログラミングで頻出のデータ構造です。巨大な配列（$N \gg 10^5$）に対するセグメント木を作ろうとするとMLEや実行時エラーを引き起こしてしまいますが、必要なノードだけ作ることで空間$O(Q)$で実現できます。これを動的セグメント木と言います。

本記事では抽象化動的セグメント木の非再帰実装の勘所を紹介します。また、効率よく実装するための工夫についても紹介します。

対象読者

本記事ではセグメント木の仕組みを理解していることを前提とします。セグメント木の解説記事は世にごまんとあるので、適宜参照してください。個人的にわかりやすいと思う記事を1つ挙げておきます。

アルゴリズム

この節ではアルゴリズムを概観します。視覚的に理解することが目的です。

一点更新

通常のセグメント木ではすべてのセグメント^[1]が適切に初期化されていますが、動的セグメント木では一点更新の度に必要に応じてセグメントをつくります。各セグメントは区間と更新されたノード、対応する区間の途中結果を持ちます。図の横軸は区間、縦軸は親子関係を表しています。

初期化

動的セグメント木は単位元で初期化されています^[2]。まだ更新していないのでセグメントはありません。

更新

初期化された動的セグメント木に1回目の更新を行います。これは、木の根をつくる操作に相当します。

2回目の更新を行います。新しいノードは根の区間の右側にあるので、右の子を生やします。セグメント木は二分木なので、$i_l < i_p < i_r$の関係を満たしている必要があります。そのため、更新するノードと根にあるノードの情報をswapします。

一点更新は行きがけにノードをswapしつつ、セグメントを追加または更新します。そして、帰りがけに途中結果を再計算します。このような処理はスタックで非再帰実装できます。

区間クエリ

区間クエリは少し複雑です。まずは、アルゴリズムを観察してみます。

クエリ区間がセグメントの片方の子に包まれるとき（Type 1）

根から順に探索していく際に、セグメントのもつノードがクエリ区間に含まれるならば答えに反映させる必要があります。二項演算の計算順序を保つために少し工夫が必要です。

セグメント木が二分木であることを思い出すと、R1の左側の子孫がもつノードのインデックスは必ず、R1のもつノードのインデックスよりも小さいです。したがって、R1のもつノードは一番右から合成することになります。他のノードも同様です。

スタックを利用して帰りがけ順序で外側から値を合成していくと、演算の順序を保つことができます。左右の区別は後述しますが、enumを使うこともできます。

区間が左右の子にまたがるとき（Type 2）

セグメントの区間とクエリ区間が完全に一致する場合は途中結果にスタックにためていた値を合成します。これが返り値です。そうでない場合、セグメントの左の子孫、自身がもつノード、右の子孫の順に値を合成します。

Pの左の子孫たちについて考えます。クエリ区間の左端はL1の左側（外側）にあります。したがって、右側（内側）の子であるL2の途中結果が答えに含まれます。L1のもつノードが区間に含まれるなら、左から合成すればよいです。

次に外側にあるL3を検討します。区間の左端が右側（内側）にあるので左側（外側）の子は不要です。L4の子孫の情報は明記されていないため不明ですが、L3のもつノードをいちばん左側（外側）から合成することは確定します。よって、L3をスタックに追加し、L4の子孫を調べていきます。

左側の子孫を調べ終わったら、Pの持つノードを必要に応じて右から合成します。右側の子孫も左側の子孫と同じように処理します。最後に、スタックに遅延していたノードの値を合成します。

外側から合成するノードはスタックに記録しておいて、帰りがけに合成できます。それ以外のノードやセグメントは順番を保って計算できます。

実装上の注意

スタックの再利用

アルゴリズムからわかるように、スタック長さは一点更新でも区間クエリでも$2\lceil \log N \rceil$で抑えることができます。動的セグメント木のフィールドに再利用するスタックを持たせておくことで、アロケーションコストを節約できます。

アリーナ木

ノード^[3]を個別にヒープに置くのではなくVecに配置するような木をアリーナ木と言います。アリーナ木はキャッシュ効率が良く、所有権の管理も楽です。

セグメントをアリーナ木に格納する際には、左右の子へのポインターをOption<NonZeroUsize>とするとよいです。Option型は通常usize1つ分の追加メモリを必要としますが、NonZeroUsizeでは0をNoneと扱う最適化が施されているため、追加メモリは必要ありません。アリーナ木の0要素目は根なので子のインデックスは常に正です。

区間クエリにおける左右の子の区別

左右の子の区別にenumを利用するとOptionと同様に追加のメモリを必要とします。アリーナ木の実態はVecですが、実はVecの容量の上限はisize::MAXです。たとえば、Vec::with_capacity()でisize::MAXより大きな容量を与えるとpanicします。したがって、usizeの最上位ビットをフラグとして利用できます。

実装例

セグメント木とその仲間たちを実装したseg_libをcrates.ioで公開しています。動的セグメント木の実装例はこちらをご覧ください。

おわりに

本記事では抽象化動的セグメント木の非再帰実装について紹介しました。一見面倒で敬遠しがちな話題ですが、図を描いてうまく整理することができました。この喜びを共有できると嬉しいです。

本記事に誤りを発見した場合、コメントにてお知らせください。また、あなたの知見を共有して頂けると嬉しいです。

脚注

1. 更新されたサイズ1のノードを「ノード」、1以上のサイズをもち途中結果が記録されるはずのノードを「セグメント」と書くことにします。「ノード」は「セグメント」ですが「セグメント」は必ずしも「ノード」ではありません。

2. 二項演算が可換な場合は、任意の値で初期化できます。動的セグメントは単位元で初期化しておき、区間クエリで更新されたノード数も計算するようにします。繰り返し自乗法で初期値のべき乗を$O(\log N)$で計算できるので、計算量は定数倍しか悪化しません。

3. ここでは普通の意味でのノード。