Segment Treeを理解したい人生だった

本記事は、Competitive Programming (1) Advent Calendar 2019 - Adventarの12/11分として書かれたものです。

この記事について

 こんにちは、monoidsigma1です。皆さんはSegment Tree(通称セグ木)を知っていますか?セグ木というのは簡単に言うと、結合則 を満たす演算の関わる区間に関するクエリを高速に処理することのできるデータ構造です。扱うことのできる構造は「モノイド」という名前が付いており、本記事でもこの言葉を使うことにします。 セグ木の導入としては区間の最小値を求める問題区間の総和を求める問題が扱われることが多いです。また、セグ木の実装を「抽象化」することで、一般のモノイドに対応することが可能であるということが広く知れ渡っています (beetさんの記事、実装についてはこちらの記事)。
 ところが、じゃあ実際セグ木でどういうことができるのか、ということが具体的に書かれた資料はあまり多くありません2。本記事では、行列演算や一次関数の合成、最大公約数の計算など有名な例はもちろん、名前のつかないような非自明なモノイド達3をセグ木で扱う問題について簡単に解説・紹介します。このようなモノイドは、あまり「モノイド感」のないモノイド達であり(理由:名前がない、有名じゃない)、彼らを扱うことができるようになれば、セグ木に対する理解をさらに深めることができると考えています。

注意

  • 問題例がいくつかあるのでネタバレを含みます。
  • 基本的に半開区間で考えます。a _ {l,r+1}とか書いてあったら区間[l,r+1)のついてのなんかの量だと思ってください。4
  • 実装例をいくつか載せていますが、セグ木の部分はうしさんのライブラリを使わせていただいています。うしさんありがとう。5

セグ木に載る構造

 セグ木をあまり知らない人(がここまで来ているかわかりませんが)のために、セグ木のアイデアを確認します。実装レベルの説明は他の記事(ふるやんさんの記事Tsutajさんの記事など)に任せることにして、どういう事実を利用しているのかということをみていきます。
 数列 a _ i が与えられていて、\min(a _ l, a _ {l+1} , \cdots, a _ r)が知りたいとしましょう。このとき、 l \leq m \leq rを満たすmについて、次の等式が成り立ちます。

\min(a _ l, a _ {l+1} , \cdots, a _ r) = \min(\min(a _ l, a _ {l+1} , \cdots, a _ {m}),\min(a _ {m+1}, a _ {l+1} , \cdots, a _ r))

これは、「区間[l,r+1)に関する最小値は、区間[l,m+1)区間[m+1,r+1)に関する最小値がわかれば計算できる」ということを言っています。再帰的な構造をしていますね。この事実が成り立つために、minについて結合則が成り立つことが重要であることに注意してください。 このよう、結合則から導かれる「ある区間の答えを、その区間を分割してできた2つの区間の結果から計算できる」という構造を利用していると考えると、セグ木を使って問題を解く上で便利だと思います。あとは計算量を落とすためにいろいろ工夫するわけですが、そこは上述の記事を参照してください。

有名な例

 それでは、セグ木を使ってできることの具体例の紹介をしていきます。

min/max,sum

 最初にセグ木を試すならやはりこれでしょう。dp高速化などの手段としても頻繁に登場します6

1次関数の合成

 f _ i=a _ i x+ b _ i (a _ i \neq 0)として、f _ r \circ f _ {r-1} \circ \cdots \circ f _ lを求める問題です。一般に関数の合成は結合則を満たしますが(交換則は満たさない!)、 次数が変わってしまうと扱うのが大変です。1次関数であれば合成しても次数は変わらないので、容易に扱うことができます。具体的には、区間[l,r+1)に対する合成の結果としてxの係数と定数項を持っておくことで、セグ木による扱いを実現できます。
yosupoさんのジャッジ Point Set Range Composite
実装例
AtCoderでの問題例 ARC008-D : タコヤキオイシクナール

行列積

 行列の積は結合則を満たすのでセグ木に乗ります。dp高速化の文脈で出てきがちな行列積ですが、行列をセグ木に乗せて解く問題もありがちです。AtCoderの問題で比較的易しいものを紹介します。

D - DISCO!

解説

まず、与えられた文字列中に部分列として"DISCO"がいくつあるか、という問題を考えます。D _ i, I _ i, S _ i, C _ i, O _ iをそれぞれ、「i文字目まで考えたときの、部分列 "D","DI","DIS","DISC","DISCO"の数」としましょう。このdpの遷移式は非常に単純ですが、この遷移を行列で書くと以下のようになります7


\begin{pmatrix}
O _ i \\\
C _ i \\\
S _ i \\\
I _ i \\\
D _ i \\\
1
\end{pmatrix}=
\begin{pmatrix}
1 & {o _ i} & 0 & 0 & 0 & 0 \\\
0 & 1 & {c _ i} & 0 & 0 & 0 \\\
0 & 0 & 1 & {s _ i}  & 0 & 0 \\\
0 & 0 & 0 & 1 & {i _ i} & 0 \\\
0 & 0 & 0 & 0 & 1 & {d _ i} \\\
0 & 0 & 0 & 0 & 0 & 1
\end{pmatrix}
\begin{pmatrix}
O _ {i-1} \\\
C _ {i-1} \\\
S _ {i-1} \\\
I _ {i-1} \\\
D _ {i-1} \\\
1
\end{pmatrix}

d _ i, i _ i, s _ i, c _ i, o_ iはそれぞれ、i文字目がD,I,S,C,Oであれば1、そうでなければ0である変数です。この行列積をとることで、dpの計算が可能です。
 元の問題に戻ります。各iに該当する行列をセグ木に乗せます。演算を通常の行列積とすれば、区間[l,r+1)に対応するdpの結果が得られることになります。行列積は可換ではないことに注意してください。
実装例

こっちの方がクエリ処理要素があります。解説はしません。 D - Dangerous Hopscotch

最大公約数(GCD)

 自然数a,bの最大公約数を求める演算gcd(a,b)結合則を満たし、セグ木に乗ります。
問題例 : C - GCD on Blackboard
もうちょっとちゃんと試したい人用(TL注意) : Codeforces Round #458 (combined) D Bash and a Tough Math Puzzle

最小共通祖先(LCA)

 2頂点のLCAはよく見ると思いますが、当然それ以上の数の頂点のLCAも定義できます。LCAを求めたい頂点集合がバラバラだと順番に試していくしかありませんが、「頂点番号が[l,r+1)の範囲にある頂点たちのLCAを求めよ」というクエリであれば、セグ木で対応可能です。区間[l,m+1)LCA区間[m+1,r+1)LCAをそれぞれ覚えておき、それら2頂点のLCAを求めれば、[l,r+1)に対するLCAを求めることができます。
問題例 : Codeforces Round #520 Div2-D Company

ここまでは演算が単純なものを紹介してきましたが、少し非自明な例を見ていきます。

ACPC2019 Day3-E 総和の切り取り

問題概要

 数列{a _ i}に対して、q回値の書き換えを行う。それぞれの書き換えの後、{a _ i}の連続部分列の総和の最大値を求めよ。
リンク

クエリでなければ有名な問題で、「dp _ i : i番目の要素を末尾とする連続部分列の総和の最大値」というdpで解けますが8、クエリがあると当然そうはいきません。セグ木を使った解法を考えてみましょう。
 ans _ {l,r+1}を数列「a _ l, a _ {l+1} , \cdots, a _ {r}における連続部分列の総和の最大値」とすれば、ans _ {0,N}を各クエリ毎に答える問題になります。「今まで通り、ansをセグ木に乗せておけばOK!」と言いたいところですが、今回はそうはいきません。次の図をご覧ください。 便宜上、区間[l,m+1)L区間[m+1,r+1)R区間[l,r+1)LRと呼ぶことにします。 f:id:sigma1113:20191212150257p:plain

この図からわかる通り、今回はans _ {LR} = \max(ans _ {L},ans _ {R})とは限りません。緑色の区間のように、左端がL内にあり、右端がRにあるような区間も考慮して、ans _ {LR}を計算しなければいけないからです。
 じゃあセグ木では解けない...?と思ってしまうかもしれませんが、そんなことはありません。ここで、「知りたい値を計算するために必要なものを一緒に載せておく」ということをします。緑色のような区間が答えとなる場合、それはどのような性質を持っているでしょうか?m+1からL側に区間を伸ばしていくことを考えると、L内において、右端をmに止めたときの区間の和が最大となるところまで伸ばしていけばいいことがわかります。R側に伸ばしていくときも同様です。従って、次の2つの量が鍵を握ることがわかります。
 sum _ {l,r+1}区間[l,r+1)の和として

  • rsum _ {l,m+1} = \max _ {l'} (sum _ {l',m+1})
  • lsum _ {m+1,r+1} = \max _ {r'} (sum _ {m+1,r'+1})

従って、これらの値を使えば、ans _ {LR}の候補は

  •  ans _ {L}
  •  ans _ {R}
  •  rsum _ {L} + lsum _ R

の3つです。これでans _ {LR}の計算を行うことができました。これで安心...と言いたいところですが、まだ考えなければいけないことがあります。今度はlsumrsumを計算しなければいけません。しかし、このためには単に区間内の和sumを持っていれば十分です。以上より、以下の4つの値を載せたセグ木を用いることで、この問題を解くことができます。

  • その区間における答えans
  • その区間内において、その区間の左端を必ず使う区間における和の最大値lsum
  • その区間内において、その区間の右端を必ず使う区間における和の最大値rsum
  • 区間の総和sum

実装については、セグ木を抽象化しておけば、「何を載せるか」ということと「区間についてどのような演算を行うか」という情報を渡してあげれば十分です。
実装例

このように、セグ木に最終的に知りたい値だけはなく、「知りたい値を計算するために必要なものを一緒に載せておく」というテクニックは非常に有用で、実際同じようにして解ける問題もAtCoder以外のサイト9では出題されています(AtCoderにもありますか?)。以下に問題例を載せておくので、興味のある人は解いてみるとよいでしょう。

略解
 載せる情報は上の問題と同じです。解答クエリは区間の位置関係によって丁寧に場合分けすることでできますが、解説によるともう少し楽にできるようです。
実装例

略解
 まず、すべての入力を読んでおいて座標圧縮をしてしまいましょう。その上で、セグ木を使って圧縮後の区間に以下の6つの情報を載せます。

  • その区間における答えans
  • 区間の左端の(本来の)座標lx
  • 区間の左端の(本来の)座標rx
  • 区間内の座標の数num
  • 区間内の各点についてのlxとの距離の総和lsum
  • 区間内の各点についてのrxとの距離の総和rsum

6つもあって少しびっくりするかもしれませんが、マージは複雑ではないです。
実装例

略解
 数学をすると、n個の鏡をクリアするのにかかる日数の期待値E _ nは以下の式で書けることがわかります。

 \displaystyle
E _ n = \sum _ {i=1}^n  \frac{1}{\Pi _ {j=i} ^ {n} p _ j}

 チェックポイントの変化に応じて、対応する値を足し引きすればよさそうです。セグ木に持たせる情報は以下の2つです。

  • その区間を抜けるのに必要な期待値(つまり答え)
  • \Pi _ {j=l} ^ {r} \frac{1}{p _ j}

実装例

略解
 4人の登場人物を0,1,2,3とします。この問題では、あるところまでは人0で、その後のあるところまでは人1で、...というように境界が定数個で、状態の変化も単調であることがポイントになります。この点を生かして、セグ木に以下の情報を持たせます。

  •  res _ {i,j} = 区間の左端を食べるのが人iで、区間の右端を食べるのが人jのときのその区間における答え

区間のマージは、各i,j (i \lt j)について、区間の境界で食べる人が切り替わる場合と、そうでない場合の2つを試せばよいです。式で書けば以下のようになります。

 \displaystyle
res _ {i,j} = \max _ {i \leq k \leq j} (\max (res _ {L _ {i,k}} + res _ {R _ {k,j}},res _ {L _ {i,k}} + res _ {R _ {k+1,j}}))

実装例

個人的にびっくりしたものを紹介しておきます。この問題を紹介するために本記事を書いたと言っても過言ではありません。ぜひ挑戦してみてください(ちなみに想定解はTLが厳しいです)。

略解
 次のものをセグ木に載せます。

  • 区間[l,r+1)の数字をすべて含むようなパスが存在するかどうか(bool値)
  • 存在するなら、そのようなパスの両端

区間のマージは、例えば各端点の組をすべて試し、パスになるかどうかを調べるという方法があります。マージにはO(1)LCAを使いましょう(は?)。また、マージの際の定数倍にも注意してください(は?)。
 この情報を持ったセグ木を使えば、Mexの値はセグ木上の二分探索によって求めることができます。

 自分はこの問題の解説を読んでびっくりしました。セグ木でこんなことができるんですね...。

 他にも「こんなものもセグ木に乗るぞ」というもので面白いものがあったら教えてください。掲載を検討します(遅延セグ木はまたの機会に...)。

RMQ and RAQ をセグ木で解く(noshi91さん提供)

 面白い話を聞いたので紹介します。RMQ and RAQは遅延セグ木の入門問題として有名ですが、なんとこれをセグ木を使って解きます。これができればもう遅延セグ木はいらないですね!
 まず簡単にアイデアを説明します。区間加算クエリをimos法で処理することを考えます。つまり、[l,r+1)x加算というクエリが来たら、lx加算し、r+1-x加算します。ここでセグ木に以下の情報を持たせます。

  • その区間より左側からの加算を無視したときの、その区間の最小値min
  • その区間より右側にあとどれくらい加算するかを表す値sum

そして区間L,Rのマージを以下のように行います。

  • sum _ {LR} = sum _ L + sum _ R
  • min _ {LR} = \min(min _ L, sum _ L + min _ R)

minのマージの式の右辺のmin _ R + sum _ L の部分がポイントです。区間Rに、区間Lからやってきたimos法の影響を足しこむことで、区間加算の操作を実現しています。
 注意点ですが、区間[l,r+1)の最小値を答えるクエリに対しては、sum _ {0,l} + min _ {l,r+1}の値を答える必要があります。lより左側から始まった区間加算クエリの影響があるかもしれないからです。

実装例

データ構造を載せる

 データ構造であるセグ木にデータ構造を載せるということもできます。CHTを載せるLi-Chao Treeセグ木を載せる話があります。本記事では紹介にとどめますが、興味のある方はリンク先の記事をご覧ください。

あとがき

 本記事では、セグ木の載るものを具体的にいくつか紹介しました。特に、知りたい値を計算するためにいろいろな値を載せるというテクニックは、セグ木の適用範囲を広めるものだと思います。本記事が皆さまのセグ木に対する理解の一助となれば幸いです。ちなみに遅延セグ木についても記事を書いていますが、公開日時は未定です。10
 それでは、快適なセグ木ライフを。


  1. idsigma - AtCoder

  2. beetくん はやくこれを完成させなさい

  3. 非自明というのは基準が人によって違うので難しいですが、代数慣れしていない人に「モノイドの例を1つ教えて」と言われて、本記事の後半で紹介するようなモノイドを挙げる人はおかしいんじゃないかと思ってしまいます(冗談です)

  4. 半開区間はいいぞ 境界を求めて累積和するときとかもごちゃごちゃにならない(個人の感想です)

  5. ライブラリを自分で書くか、ネットに落ちているものを使うか、それとも用意しないかというところは意見がわかれるところだと思います。中身がブラックボックスだと困ることもあるので、最初に一度は自分で書いて、あとは早いのを使うというのが個人的な落としどころです(実践しているとは言っていない)

  6. 実家dpとか誰かが言ってたらだいたいセグ木が出てくるイメージ

  7. 1いらないかも

  8. i番目を足しても負にならないなら足しこむ、なるなら0を入れる感じの遷移

  9. こういうのも面白いと思うんですけどね…

  10. 説明がつらい