0. はじめに
2020年9月7日にAtCoder公式のアルゴリズム集 AtCoder Library (ACL)が公開されました。
私はACLに収録されているアルゴリズムのほとんどが初見だったのでいい機会だと思い、アルゴリズムの勉強からPythonでの実装までを行いました。
この記事ではSegtree(セグメント木)をみていきます。
対象としている読者
- セグメント木ってなに？という方。
- セグメント木を理解し自分で実装できるようにしたい方。
- ACLのコードを見てみたけど何をしているのかわからない方。
- C++はわからないのでPythonで読み進めたい方。
参考にしたもの
関連する数学の参考記事です。
maspyさんのセグメント木に関する記事です。$N \ne 2^n$ の場合のセグメント木についても図示されています。
1. セグメント木とは
セグメント木とは完全二分木の構造を利用することで列(「数」列や、「文字列」列など)に対して以下の2種類のクエリ
- $i, x$ が与えられたとき、$i$ 番目の要素を $x$ に更新（更新クエリ）
- $l, r$ が与えられたとき、区間 $[l, r)$ の要素の総積を取得（取得クエリ）
を高速に行うことができるデータ構造です。（「積」は掛け算に限らず演算の結果を表すこととします。）
もしこの段階で、「クエリってなに？」、「なんで木を使うと高速化できるの？」、「何言ってるか全然わからん」と感じた場合は先にFenwick_Tree編（の１〜２章）を読んでから戻ってくると良いかもしれません。
1.1. ACLのセグメント木
ACLのセグメント木は非再帰・抽象化セグメント木です。
- 非再帰：再帰呼出しを用いずに実装します。
- 抽象化：要素間の演算を固定せず、使用時に指定できるように実装します。
これに倣い、本記事では具体的な演算に特定せず、演算は演算として扱うこととします。(理解のために具体的な演算で例示することはします。)
1.2. 本記事における用語、記号の定義
まず、本記事における用語、記号の意味を確認しておきます。
集合
クエリの対象となる列 $A$ の要素 $a_0, a_1, \cdots , a_{N-1}$ が属する集合を $S$ と呼ぶことにします。
例)
- $A = [1, -2, 3] \rightarrow Sは整数全体$
- $A = ['a', 'b', 'c'] \rightarrow Sは英小文字アルファベット全体$
- $A = [('alice', -11.7), ('bob', 12.5)] \rightarrow Sは(英小文字の文字列, 実数)というペア全体$
もちろん１番目の例で $S$ は整数全体を包含する集合である実数全体などといっても良いです。
二項演算
集合 $S$ の任意の要素２つ(同じ要素でも良い)から新たな値を得る規則を $\boldsymbol{S}$ 上の二項演算といいます。本記事ではこれを単に「演算」と呼び、特に要素 $x, y$ に対する演算を $op(x, y)$ と書き、演算子として記号 $\star$ を用いることにします。すなわち、
op(x, y) = x \star y
です。
例１ $S$ を自然数全体とします。このとき
-
演算として足し算を考えると $S$ の要素 $2$ と $3$ の積は $op(2, 3) = 2 \star 3 = 5$ になります。
-
$op(x, y) := (文字 'a' をx個、文字 'b' を y 個をこの順につなげる)$ というものを考えると $S$ の要素 $2$ と $3$ の積は $op(2, 3) =; 'aabbb'$ になります。
-
「文字列 $x$ の2文字目と文字列 $y$ の5文字目をこの順につなげる」という規則は $x=2$ 等の不適な要素が $S$ に存在するので $S$ 上の演算とは言えません。
例２ $S$ を空でない文字列全体とします。このとき
-
$op(x, y):= (xの1文字目とyの1文字目をこの順につなげる)$ とすると $S$ の要素 $'bob'$ と $'alice'$ の積 は $op('bob', ~'alice') =~ 'bob' \star ~'alice' =~ 'ba'$ となります。
-
$op(x, y):= (xの長さ \div yの長さ)$ とすると $S$ の要素 $'bob'$ と $'alice'$ の積 は $op('bob', ~'alice') = 0.6$ となります。
例１の１つ目で見た、「自然数全体上での足し算」は積も自然数全体という集合に属します。また、例２の１つ目で見た、「空でない文字列全体上での各1文字目の連結」という演算も、積は空でない文字列全体という集合に属します。このように、「積全体が属する集合」が「演算が定義された集合」に含まれる場合、演算$op$は集合$S$において閉じているといいます。
一方、例1, 2の２つ目の演算は閉じていません。演算が閉じていないと扱いが難しいので本記事では閉じた演算を考えることにします。もっとも、整数全体の集合上での足し算、掛け算、最大値、最小値などは閉じているので、あまり気にしなくても良いかもしれません。
2. セグメント木の仕組み
ではセグメント木の仕組みを見ていきます。
ここからはクエリの対象となる列として長さ $N=7$ の列 $(a_0, a_1, \cdots, a_{6})$ を考えます。
また、前章で述べたように、この列の要素が属する集合を $S$ 、$S$ 上の演算を $op$ と呼び、その演算子を $\star$ とします。また、演算 $op$ は集合 $S$ において閉じているとします。
2.1. 完全二分木
1.セグメント木とはで書いたように、セグメント木は完全二分木の木構造です。すなわち、葉以外の全てのノードが２つの子ノードを持ち、全ての葉は同じ深さを持ちます。
いま、$N=7$ なのでこの列の各要素を葉に分配するために葉の数が $8$ の完全二分木を用意します。
葉以外の各ノードには子ノードの積を格納します。つまり、子ノードが $a_0$ と $a_1$ のノードには $a_0 \star a_1$ を格納し、子ノードが $a_0 \star a_1$ と $a_2 \star a_3$ のノードには $(a_0 \star a_1) \star (a_2 \star a_3)$ を格納します。
さて、ここで２つの問題が現れます。
１つ目は上図の上から２段目に現れるかっこについてです。いま、取得クエリとして $l = 0, r = 4$ の場合を考えます。すなわち、区間 $[0, 4)$ の要素の総積
a_0 \star a_1 \star a_2 \star a_3
を取得することを考えます。上図をみてみるとちょうど $(a_0 \star a_1) \star (a_2 \star a_3)$ が格納されたノードが存在します。しかし、問題は
(a_0 \star a_1) \star (a_2 \star a_3) = a_0 \star a_1 \star a_2 \star a_3
は必ず成り立つとは限らないということです。
例えば、要素として $(a_0, a_1, a_2, a_3) = (0, 1, 2, 3)$、演算として $op(x, y) := x + 2y$ というものを考えてみます。
左辺はかっこの部分を先に計算し、
\begin{aligned}
(a_0 \star a_1) \star (a_2 \star a_3)
=& 2 \star 8\\
=& 18
\end{aligned}
です。一方右辺は左から順番に計算し、
\begin{aligned}
a_0 \star a_1 \star a_2 \star a_3
=& 2 \star 2 \star 3\\
=& 6 \star 3\\
=& 12
\end{aligned}
となり、等号は成り立ちません。
一方、演算として足し算 $op(x, y) := x + y$ を考えると、知っての通り等号が成り立ちます。
このように、演算によっては計算する順序が重要になる場合があります。ですが、計算する順序によって値が異なる演算の場合にはセグメント木の「事前に２つずつの積を保持しておく」という操作が無意味になってしまうので計算する順序を変えても値が変わらない演算のみを考えることにします。
具体的には任意の要素 $x, y, z \in S$ について
(x \star y) \star z = x \star (y \star z)
という関係が成り立つ演算のみを考えます。この関係を結合法則(結合律)といいます。
注意 交換法則について
本記事の実装では演算が交換法則を満たすことを要求しません。
交換法則とは
x \star y = y \star x
という関係です。これを要求しないということは $x \star y$ と $y \star x$ を区別するということです。（演算として引き算などを考えれば交換法則が常に成り立つわけではないことがわかると思います。）
２つ目の問題は上図の"？"の部分です。完全二分木の葉の数は 2 のべき乗個ですが、対象の列の要素が常に 2 のべき乗個とは限りません。空いている部分には何を入れれば良いでしょうか？
ここでも、取得クエリを考えてみます。
- $l = 2, r = 6$ のとき
$a_2 \star a_3$ というノードと $a_4 \star a_5$ というノードの積を取れば良いです。 - $l = 4, r = 7$ のとき
$a_4 \star a_5$ というノードと $a_6$ というノードの積を取れば良いです。 - $l = 0, r = 7$ のとき
$a_0 \star a_1 \star a_2 \star a_3$ というノードと $a_4 \star a_5$ というノードと $a_6$ というノードの積を取れば良いです。（先ほど結合法則が成り立つ演算のみを考えることにしたのでかっこはつきません。）
実は、? の部分は使わずに取得クエリを処理することができます。なので、? には $S$ の要素であれば何を入れても問題なく動作します。（$S$ に属さない要素を入れると演算が行えないので困ります。）
しかし、ACLのセグメント木ではある特別な値を入れています。この特別な値を用意することで
- 全要素の総積が $O(1)$ で取得できる
- セグメント木の初期化方法が柔軟になる
- 実装が容易になる
といったメリットがあります。
ここでは「全要素の総積を $O(1)$ で取得するため」という理由付けをしてどのような値を入れれば良いかみていきます。
全要素の総積が $O(1)$ で取得できるということはすでに計算された結果がどこかのノードに格納されているということです。もしそのような値が入っているとしたらそれは根の部分でしょう。つまり図の最上段のノードには
a_0 \star a_1 \star a_2 \star a_3 \star a_4 \star a_5 \star a_6
が入っていることになります。ここから下段の ? を逆算していきます。まず、１つ下の段は左側が
a_0 \star a_1 \star a_2 \star a_3
なので右側の ? には
a_4 \star a_5 \star a_6
が入ることがわかります。さらにこの下を見ると次の ? は
a_6
であることがわかります。さて、これはどういうことでしょうか。この $a_6$ というのは左の子ノードと一致しています。式で書けば
a_6 \star ? = a_6
ということです。いま、$a_6$ は $S$ の要素であり、$S$ 内のあらゆる値を取り得ます。つまり、? は $\boldsymbol{S}$ の任意の要素に対して演算をしても影響を与えないような値です。このような特別な値を慣例に倣って $e$ と書くことにします。
具体的には、葉ノードの ? には次のような値 $e$ を入れます:
任意の $x \in S$ に対して
x \star e = e \star x = x
を満たす。
このような $e$ を単位元といいます。
例をみてみましょう。$S$ を整数全体とします。
- $op(x, y) = x + y$ のとき
任意の要素 $x \in S$ に対して $x + 0 = 0 + x = x$ なので単位元は $0$ です。 - $op(x, y) = x \cdot y$ のとき
任意の要素 $x \in S$ に対して $x \cdot 1 = 1 \cdot x = x$ なので単位元は $1$ です。
このように単位元は集合 $S$ や演算の種類によって異なります。また、場合によっては単位元が存在しないこともあります。
ここまでをまとめると、（ACLの実装で）セグメント木を構築するには次の条件を満たしている必要があることがわかりました。
- 演算 $op$ が集合 $S$ において閉じている
- $op(x, y) = x \star y$ は結合法則 $(x \star y) \star z = x \star (y \star z)$ を満たす
- $S$ に単位元 $e$ が存在する
これらの条件を満たす組 $(S, op, e)$ をモノイドといいます。
ここからはこれらの条件を満たしている場合を考えます。
するとセグメント木の各ノードは下図のようになります。
2.2. 木を１次元配列で表現する
次に、木を１次元配列で表現することを考えます。木構造で重要なのはノードの親子関係が決まっていることです。言い換えれば、あるノードからそのノードの親や子にアクセスさえできれば、形として木になっている必要はありません。セグメント木は完全二分木という整った構造のため、適切に各ノードにindexを振ることで親子のアクセスを容易にすることができます。
具体的にはdataという長さ $2 \times (葉の数)$ の１次元配列を用意し、下図（赤字）のようにindexを対応させます。また、葉の数はこの後何度か出てくるので便宜上これを size と呼ぶことにします。本記事の例では size=8 です。
では、１次元配列dataに対するアクセスをみていきます。
まず、クエリ対象の列の要素にアクセスする方法です。例えば、$a_0$ を取得する場合には data[8]、$a_3$ を取得する場合にはdata[11]とアクセスすれば良いです。より一般的には、$a_i$ を取得するにはdata[i + size]へのアクセスをすれば良いことがわかります。
次に子ノードへアクセスする方法です。いくつかみてみると、
-
data[1]の子はdata[2]とdata[3] -
data[2]の子はdata[4]とdata[5] -
data[5]の子はdata[10]とdata[11]
となっています。一般的には
-
data[i]の子はdata[2 * i]とdata[2 * i + 1]
であることがわかります。
最後に親ノードへのアクセスです。
-
data[8]の親はdata[4] -
data[9]の親はdata[4] -
data[6]の親はdata[3] -
data[7]の親はdata[3]
となっていて、一般的に
-
data[i]の親はdata[i // 2]
であることがわかります。
これで親子へのアクセス方法がわかったので木を１次元配列で表現することができました。
ここからは構築、クエリ処理の具体的な方法をみていきます。
2.3. セグメント木の構築
ではセグメント木の構築方法からみていきます。
手順としては
- セグメント木の大きさを決める
- 葉に対象の列を入れる
- 葉以外のノードに区間の演算結果を入れる
となっています。
擬似コードは以下のようになります。配列のスライスアクセス[l:r]は右開区間[l, r)です。
セグメント木の構築 (入力: 列A, 演算op, 単位元e){
N = Aの長さ
size = N以上で２のべき乗となる最小の数
data = 長さ(2 * size)、初期値eの配列
// 葉に入力の列を格納
data[size:size + N] = A
// 葉以外のノードを更新
for i = ((size - 1) to 1){
data[i] = op(data[i]の左子ノード, data[i]の右子ノード)
}
}
葉以外のノードを更新する際には
- 葉に近い部分から更新する必要がある
- 演算の左右を入れ替えてはいけない
という点に注意が必要です。
2.4. 更新クエリ
更新クエリ: $i, x$ が与えられたとき、要素 $a_i$ を $x$ に更新
を処理する方法を考えます。
いくつか具体例をみてみます。
① $i = 2, x = b_2$ のとき
まず、$a_2$ を $b_2$ に更新します。これはdata[2 + size] = b_2とすればいいです。
ここから、$a_2$ が関係する場所を全て更新していきます。そのためには
data[5] = op(data[10], data[11])data[2] = op(data[4], data[5])data[1] = op(data[2], data[3])
とすればいいです。
② $i = 6, x = b_6$ のとき
まず、data[6 + size] = b_6とします。
次に $a_6$ が格納されているdata[7]を更新しますが、data[7]の中身は実際には $a_6 \star e$ だということを思い出せば先ほどの例と同様に
data[7] = op(data[14], data[15])
とすればいいことがわかります。その後
data[3] = op(data[6], data[7])data[1] = op(data[2], data[3])
と更新します。
より一般的には以下のようになります。
更新クエリ(入力： i, x){
// 葉に移動し更新
i += size
data[i] = x
//関係する場所を更新
while (True){
iをdata[i]の親のindexに変更
data[i] = op(data[i]の左子ノード, data[i]の右子ノード)
もしiが根(=1)なら終了
}
}
2.5. 取得クエリ
取得クエリ: $l, r$ が与えられたとき、区間 $[l, r)$ の要素の総積を取得
を処理する方法を考えます。
いくつか具体例をみてみます。
① $l = 0, r = 3$ のとき
区間が右開区間であることに注意します。
a_0 \star a_1 \star a_2 = (a_0 \star a_1) \star a_2
なので $op(\rm{data}[4], \rm{data}[10])$ が求めるものです。
② $l = 1, r = 7$ のとき
\begin{aligned}
&a_1 \star a_2 \star a_3 \star a_4 \star a_5 \star a_6\\
=& a_1 \star (a_2 \star a_3) \star (a_4 \star a_5) \star a_6\\
=& \rm{data}[9] \star \rm{data}[5] \star \rm{data}[6] \star \rm{data}[14]
\end{aligned}
なので４ヶ所を見れば良いことがわかります。
ではこのような操作を機械的に行うにはどうすれば良いでしょうか。様々な例をみてみると次のことがわかります：
「図の各段においてみるべき場所は、区間 $[l, r)$ に完全に含まれるノードの内、両端の高々２ヶ所である。」
なぜなら、区間 $[l, r)$ は連続した区間であり、もしある段で３ヶ所みるべき場所があった場合、そのうちの２ヶ所は必ず親が共通しているからです。また、結合法則が成り立っているので、左から順番に計算する必要はありません。（交換法則は要求しないので左右の入れ替えはNGです。）そこで、下図のように「左の結果」、「右の結果」という２つの変数を用意し、区間 $[l, r)$ の両側から各段についてノードの値を使うべきかを判断し、最後に２つをまとめることを考えます。
擬似コードにすると次のようになります。右開区間であること、演算の左右は入れ替えられないことに注意します。
取得クエリ(入力: l, r){
// 最下段(葉ノード)に移動
l += size
r += size
左の結果 = e(単位元)
右の結果 = e
while (未計算の区間がなくなるまで){
data[l]を使う場合:
左の結果 = op(左の結果, data[l])
data[r-1]を使う場合:
右の結果 = op(data[r-1], 右の結果)
l, rを一つ上の段の適切なノードに移動
}
return op(左の結果, 右の結果)
}
ここからはこの方針に基づき、左右それぞれの場合に、
- ノードの値を使う条件
- 次にみるノードへの移動の仕方
をみていきます。
１. 左側の場合
ノードを使うかどうかは「今いるノードが親ノードにとって左子ノードか右子ノードか」によって判断できます。
いま、data[l]は左子ノードだったとします。このとき、data[l]の親は
data[l]-
data[l]より右の要素であるdata[l+1]
を含んでいるので、data[l]を使うメリットはありません。（data[l]の親を使うことでdata[l]を使うよりも広い区間を一度に求めることができるからです。）
一方、data[l]が右子ノードだった場合はどうでしょうか。このとき、data[l]の親は
data[l]-
data[l]より左の要素であるdata[l-1]
を含んでいます。したがって、data[l]の親を使うことはできないので、data[l]は使わなければならないことがわかります。
また、次の段への遷移としては
- 左子ノードの場合 $\rightarrow$
data[l]の親へ移動 - 右子ノードの場合 $\rightarrow$
data[l+1]の親へ移動
とすれば良いことがわかります。
2. 右側の場合
右側も考え方は左側と同じですが、区間 $[l, r)$ は右開区間であることに注意が必要です。つまり、$r$ に対してみるべき場所は常に $r - 1$ となります。
dara[r]が右子ノードの場合、区間の右端はdata[r-1]であり、data[r-1]の親は
data[r-1]-
data[r-1]より右の要素であるdata[r]
を含むのでdata[r-1]を使う必要があります。
一方data[r]が左子ノードの場合、data[r-1]の親はdata[r-1]より広範囲を一度に求められるのでdata[r-1]を使うメリットはありません。
また、次の段への遷移としては
- 左子ノードの場合 $\rightarrow$
data[r]の親へ移動 - 右子ノードの場合 $\rightarrow$
data[r-1]の親へ移動
とすれば良いことがわかります。（$r$ は右端が開区間となるように移動します。）
これを未計算の区間がなくなるまで繰り返します。つまり、終了条件は $l = r$ です。
以上をまとめると取得クエリの擬似コードは次のようになります。
取得クエリ(入力: l, r){
// 最下段(葉ノード)に移動
l += size
r += size
左の結果 = e(単位元)
右の結果 = e
while (l < r){
data[l]が右子ノードの場合:
左の結果 = op(左の結果, data[l])
l -> l + 1 に移動
data[r]が右子ノードの場合:
r -> r - 1に移動
右の結果 = op(data[r], 右の結果)
l, rをそれぞれ親に移動
}
return op(左の結果, 右の結果)
}
実際にこのアルゴリズムが正しく動作することを確認しておきます。
例1. $l = 1, r = 5$ の場合
- 葉に移動: $l = 9, r = 13$
-
data[9]は右子ノードなので使用し、lに $1$ を加算: $左の結果 = a_1, l = 10$ -
data[13]は右子ノードなのでrから $1$ を減算し、data[12]を使用: $r = 12, 右の結果 = a_4$ - l, rを親に移動: $l = 5, r = 6$
-
data[5]は右子ノードなので使用し、lに $1$ を加算: $左の結果 = a_1 \star a_2 \star a_3 , l = 6$ -
data[6]は左子ノードなので使用しない - l, rを親に移動: $l = 3, r = 3$
- $l = r$ なので左の結果と右の結果を演算し終了:
return$a_1 \star a_2 \star a_3 \star a_4$
例2. $l = 2, r = 3$ の場合
-
葉に移動: $l = 10, r = 11$
-
data[10]は左子ノードなので使用しない -
data[11]は右子ノードなのでrから $1$ を減算し、data[10]を使用: $右の結果 = a_2$ -
l, rを親に移動: $l = 5, r = 5$
-
$l = r$ なので左の結果と右の結果を演算し終了: $左の結果 = e, 右の結果 = a_2$ ,
return$a_2$
3. セグメント木の実装
ではPythonで実装していきます。変数名、関数名は出来るだけACLに合わせて実装します。
また、各メソッドについて時間計算量を記します。ただし、クエリ対象の列の要素数を $N$ とし、与える演算 $op$ の時間計算量を $f_{op}(N)$ とします。（単位元は $O(1)$ とします。)
3.1. コンストラクタ
まず、クラスSegTreeを作成し、コンストラクタを実装します。コンストラクタは2.3. セグメント木の構築に相当します。
引数は演算op、単位元e、対象の列の要素数n、対象の列vです。e、opは関数オブジェクトです。vはデフォルトでNoneとしており、vを与えない場合は全てのノードが単位元のセグメント木が構築されます。_dがこれまでの説明の $data$ に相当します。
ACLでは_log（$2^x \geq n$ となる最小の整数 $x$）を求める関数としてatcoder/internal_bitにceil_pow2()というものを定義していますが、Pythonではbit_length()メソッドを使うことで求められるのでこれで代用します。
例：$n = 7$の場合
\begin{aligned}
(n - 1).\rm{bit\_length()} &= \rm{bit\_length(6)}\\
&= \rm{bit\_length((110)_2)}\\
&= 3
\end{aligned}
class SegTree:
def __init__(self, op, e, n, v=None):
self._n = n
self._op = op
self._e = e
self._log = (n - 1).bit_length() # 2^(_log) >= n となる最小の整数
self._size = 1 << self._log
self._d = [self._e()] * (2 * self._size)
if v is not None:
# 葉に対象の列を格納
for i in range(self._n):
self._d[self._size + i] = v[i]
# 葉に近い場所から順に更新
for i in range(self._size - 1, 0, -1):
self._update(i)
def _update(self, k):
self._d[k] = self._op(self._d[2 * k], self._d[2 * k + 1])
時間計算量は
-
vを与える場合、$O(f_{op}(N)\cdot N)$ -
vを与えない場合、$O(N)$
です。
3.2. 更新クエリ
更新クエリ: $a_p$ を $x$ に更新
を実装します。
def set(self, p, x):
assert 0 <= p < self._n
# 葉に移動
p += self._size
self._d[p] = x
# 関連する場所を更新
for i in range(1, self._log + 1):
self._update(p >> i)
更新する場所は木の高さ分あるので時間計算量は $O(f_{op}(N)\log(N))$ です。
3.3. 取得クエリ
取得クエリ: 区間 $[l, r)$ の総積を取得
を実装します。自身が右子ノードかどうかはindexの偶奇で判断できます。本実装ではindexが奇数の場合自身は右子ノードとなるので、 $1$ とのandを取ることで判別できます。
$l = r$ の場合、単位元が返ります。
def prod(self, l, r):
assert 0 <= l <= r <= self._n
# 左の結果、右の結果
sml, smr = self._e(), self._e()
l += self._size
r += self._size
# 未計算の区間がなくなるまで
while l < r:
# 自身が右子ノードなら使用
if l & 1:
sml = self._op(sml, self._d[l])
l += 1
if r & 1:
r -= 1
smr = self._op(self._d[r], smr)
# 親に移動
l >>= 1
r >>= 1
return self._op(sml, smr)
みる場所は木の高さ分あるので時間計算量は $O(f_{op}(N)\log(N))$ です。
3.4. その他のメソッド
全要素の総積を取得するメソッドを実装します。3.1. 完全二分木で述べたように、全要素の総積は根に入っているのでこれを取得すれば良いです。
def all_prod(self):
return self._d[1]
時間計算量は $O(1)$ です。
$p$ が与えられたとき、$a_p$ を取得するメソッドを実装します。$a_p$ は葉に入っているのでこれを取得すれば良いです。
def get(self, p):
assert 0 <= p < self._n
return self._d[p + self._size]
時間計算量は $O(1)$ です。
3.5. まとめ
まとめたものを載せておきます。
class SegTree:
def __init__(self, op, e, n, v=None):
self._n = n
self._op = op
self._e = e
self._log = (n - 1).bit_length()
self._size = 1 << self._log
self._d = [self._e()] * (2 * self._size)
if v is not None:
for i in range(self._n):
self._d[self._size + i] = v[i]
for i in range(self._size - 1, 0, -1):
self._update(i)
def set(self, p, x):
assert 0 <= p < self._n
p += self._size
self._d[p] = x
for i in range(1, self._log + 1):
self._update(p >> i)
def get(self, p):
assert 0 <= p < self._n
return self._d[p + self._size]
def prod(self, l, r):
assert 0 <= l <= r <= self._n
sml, smr = self._e(), self._e()
l += self._size
r += self._size
while l < r:
if l & 1:
sml = self._op(sml, self._d[l])
l += 1
if r & 1:
r -= 1
smr = self._op(self._d[r], smr)
l >>= 1
r >>= 1
return self._op(sml, smr)
def all_prod(self):
return self._d[1]
def _update(self, k):
self._d[k] = self._op(self._d[2 * k], self._d[2 * k + 1])
3.6. 実行速度の改善
ACLはC++で書かれており、コンパイラによる最適化がなされることが前提になっているため、アルゴリズムの記述の仕方も真似して書くとPythonのようなインタープリタ言語では実行速度が想定以上に遅くなってしまいます。
例えばABC185-F Range Xor Queryは先ほど実装したセグメント木を用いて解くことができますが Python(3.8.2)での提出はTLEでした。
解決手段としては
- PyPyやNumbaを用いてコンパイルする
- 遅くなる原因を探し改善を頑張る
があると思います。特にPyPyで実行することによるコンパイルは手軽であり、多くの人が使っている手法だと思います。
ここでは以下の２つを(汎用性を損なわずに)改善し ABC185-F Range Xor Query をPython(3.8.2)でACできるようにします。
① 関数呼び出しを減らす
Pythonでの関数の呼び出しはコストが大きいのでこれを減らします。先ほどの実装ではコンストラクタとsetメソッドで共通して行う処理として_updateメソッドを作りましたがこれをコンストラクタおよびsetメソッド内に直接書くことにします。
② assert文をなくす
ACLではassert文が多用されていますがPythonの場合AtCoder上での実行時にassertの条件式が（おそらく）評価されてしまうのでこれをなくすことで若干の改善ができると思います。
これらを行うと以下のようになります。
class SegTree:
def __init__(self, op, e, n, v=None):
self._n = n
self._op = op
self._e = e
self._log = (n - 1).bit_length()
self._size = 1 << self._log
self._d = [self._e()] * (self._size << 1)
if v is not None:
for i in range(self._n):
self._d[self._size + i] = v[i]
for i in range(self._size - 1, 0, -1):
self._d[i] = self._op(self._d[i << 1], self._d[i << 1 | 1])
def set(self, p, x):
p += self._size
self._d[p] = x
while p:
self._d[p >> 1] = self._op(self._d[p], self._d[p ^ 1])
p >>= 1
def get(self, p):
return self._d[p + self._size]
def prod(self, l, r):
sml, smr = self._e(), self._e()
l += self._size
r += self._size
while l < r:
if l & 1:
sml = self._op(sml, self._d[l])
l += 1
if r & 1:
r -= 1
smr = self._op(self._d[r], smr)
l >>= 1
r >>= 1
return self._op(sml, smr)
def all_prod(self):
return self._d[1]
この実装でABC185-F Range Xor Queryに提出してみると 2912ms でACすることができました。
元の実装では 3300msだったので400ms（あるいはそれ以上）の改善ができました。
もちろんPyPyやNumbaを使えばより高速化でき1000ms未満でACできます。
4. 使用例
ABC185-F Range Xor Queryを実装したセグメント木(実行速度改善ver)を使って解いてみます。
問題文
長さ $N$ の整数列 $A$ があります。
あなたは今からこの数列について $Q$ 個のクエリを処理します。
$i$ 番目のクエリでは、 $T_i, X_i, Y_i$ が与えられるので、以下の処理をしてください。
- $T_i=1$ のとき
$A_{X_i}$ を $A_{X_i}⊕Y_i$ で置き換える- $T_i=2$ のとき
$A_{X_i}⊕A_{X_{i+1}}⊕A_{X_{i+2}}⊕\cdots⊕A_{Y_i}$ を出力する
ただし $a⊕b$ は $a$ と $b$ のビット単位 xor を表します。
$T = 1$ は指定された場所を指定された値で更新する更新クエリ、$T = 2$ は指定された区間の総積を取得する取得クエリになっているのでセグメント木を用いて解けそうです。
では実際に解いていきます。
まず、セグメント木の使用条件を確認します。
条件は
- 演算 $op$ が集合 $S$ において閉じている
- $op(x, y) = x \star y$ は結合法則 $(x \star y) \star z = x \star (y \star z)$ を満たす
- $S$ に単位元 $e$ が存在する
でした。
この問題では制約から $A_i$ が $0$ 以上の整数であることが保証されているので集合 $S$ を $0$ 以上の整数全体とします。そして演算はもちろんビット毎のxorなので
op(x, y) := x ⊕ y
です。Pythonでの対応する演算子は ^ です。また、この演算における単位元は $0$ であり、これは $S$ の要素として存在しています。
したがって、
- 集合 $S$ : $0$ 以上の整数全体
- 演算 $op$ : ビット毎のxor
- 単位元 $e$ : $0$
とすると全ての条件を満たしていることが確認できます。（各自確認してください。）
よってセグメント木が使えることがわかりました。
セグメント木のコンストラクタは引数に演算、単位元を関数オブジェクトとして与える必要があるのでこれを作成しセグメント木を構築します。
def op(x, y):
return x ^ y
def e():
return 0
seg = SegTree(op, e, n, A) # nはAの長さ, Aはリスト
あとは各クエリを処理するコードを書けば終わりです。
$T = 1$ のときは $A_{X_i}$ を $A_{X_i}⊕Y_i$ に更新するのでまずgetメソッドで $A_{X_i}$ を取得しこれを $Y_i$ とのxorをとりsetメソッドで更新します。
$T=2$ のときはprodメソッドが右開区間であることに注意してprodを行います。
t, x, y = map(int, input().split())
x -= 1
if t == 1:
seg.set(x, seg.get(x) ^ y)
else:
print(seg.prod(x, y))
これでこの問題を解くことができました。
提出例
5. セグメント木の補足
5.1. 単位元が存在しないとき
状況によっては集合 $S$ に演算 $op$ の単位元 $e$ が存在しないことがあるかもしれません。実は単位元が存在しなくても以下の手順を踏むことで本記事で実装したセグメント木を使用することができます。
- $e^\prime \not\in S$ である要素 $e^\prime$ を用意する
- $S^\prime$ を $S$ に $e^\prime$ を加えた集合とする
- $S^\prime$ 上の演算として $op^\prime$ を $e^\prime$ が単位元となるように適当に定義する
こうしてできた $(S^\prime, op^\prime, e^\prime)$ はセグメント木の使用条件を満たします。(ただし、元の演算 $op$ は $S$ 上で閉じていて、かつ結合法則を満たすことを仮定しています。)
例
$S$ を長さ1以上の文字列全体の集合、$op$ を文字列の連結とします。このとき $S$ には単位元が存在しません。（空文字 '' は $S$ の要素ではありません。）
そこで
- $S$ に属さない要素である整数 $0$ を用意しこれを $e^\prime$ とする
- $S^\prime$ を「長さ1以上の文字列全体の集合に $e^\prime (= 0)$ を加えた集合」とする
- $S^\prime$ 上の演算として $op^\prime$ を以下のように定義する
def op_prime(a, b):
if a == 0 :
return b
elif b == 0:
return a
else:
return a + b
とすることでセグメント木を使用できます。
(もちろん $e^\prime$ として空文字を用意すればより簡潔に書けますし返り値の型も揃いますが、ここでは単位元はどんなものでもいいということを強調するためにあえて整数の $0$ を使いました。)
5.2. 完全二分木にする必要性
本記事の実装では葉の数を２のべき乗として、空いている部分には単位元を入れる実装をしましたが、実はクエリ対象となる長さ $N$ の列 $A$ に対し、長さ $2N$ の配列 data を用意し data[N:] = A とすれば 同じコードで正常に動作します。
ただし、演算が非可換の場合には根には全要素の（正しい順序での）総積が入らないといった違いがあります。
詳しくはmaspyさんのページにありますのでそちらを参考にしてください。
6. セグメント木上の二分探索
セグメント木自体は実装できましたが、ACLのセグメント木には他にもmax_right、min_leftというメソッドが用意されています。ここからはこれらのメソッドをみていきます。
6.1. 新たなクエリ
先ほどまでの更新クエリ、取得クエリに加えて以下の新たなクエリを考えます。
以降、$prod(l, r) ;(l \leq r)$ は区間 $[l, r)$ の要素の総積を表し、特に、 $l = r$ の時は $prod(l, r) = e$ （単位元）とします。
右探索クエリ
$0 \leq l < N$ なる整数 $l$ および条件式 $f$ が与えられたとき、$f(prod(l, r)) = True$ となる最大の $r$ を求める。ただし、与えられる条件式 $f$ は次を満たす:
ある整数 $x > l$ に対し、$f(prod(l, x))=True$ であるとき、任意の整数 $l < y \leq x$ について $f(prod(l, y))=True$ である。また、$f(e) = True$ である。
このクエリを $max\_right(l, f)$ と書く。
左探索クエリ
$0 < r \leq N$ なる整数 $r$ および条件式 $f$ が与えられたとき、$f(prod(l, r)) = True$ となる最小の $l$ を求める。ただし、与えられる条件式 $f$ は次を満たす:
ある整数 $x < r$ に対し、$f(prod(x, r))=True$ であるとき、任意の整数 $x \leq y < r$ について $f(prod(y, r))=True$ である。また、$f(e) = True$ である。
このクエリを $min\_left(r, f)$ と書く。
右探索クエリと左探索クエリは向きが違うだけなのでここからは右探索クエリのみを考えることにします。
まずは具体例をみてみます。
クエリ対象の列 $A$ を
A = [0, 1, 2, 3, 4, 5, 6]
とし、演算 $op$ を足し算とします。
① $l = 0, f(x) := bool(x < 5)$ のとき
- $prod(0, 1) = 0$ なので $f(prod(0, 1))=True$ です。
- $prod(0, 2) = 1$ なので $f(prod(0, 2))=True$ です。
- $prod(0, 3) = 3$ なので $f(prod(0, 3))=True$ です。
- $prod(0, 4) = 6$ なので $f(prod(0, 4))=False$ です。
したがって、$max\_right(0, x < 5) = 3$ です。
② $l = 4, f(x) := bool(x \geq 4)$ のとき
- $prod(4, 5) = 4$ なので $f(prod(4, 5))=True$ です。
- $prod(4, 6) = 9$ なので $f(prod(4, 6))=True$ です。
- $prod(4, 7) = 15$ なので $f(prod(4, 7))=True$ です。
したがって、$max\_right(4, x \geq 4) = 7$ です。
③ $l = 3, f(x) := bool(x \leq 2)$ のとき
- $prod(3, 4) = 3$ なので $f(prod(3, 4))=False$ です。
したがって、$max\_right(3, x \leq 2) = 3$ です。(左端が $3$ であり条件式を満たす最大の区間は $[3, 3)$、 つまりそのような長さが正の区間は存在しません。)
6.2. 探索クエリの時間計算量
このクエリの時間計算量を考えてみます。なお、ここでは条件式 $f$ の評価は $O(1)$ でできるものとします。
まず、ナイーブな解法としては先ほどの具体例でみたように、$r = l + 1$ から順番にみていき、$f(prod(l, r)) = False$ となったときの $r-1$ が答えです。$prod(l, r)$ はセグメント木を用いることで $O(f_{op}(N)\log(N))$ で取得できるので、全体では $O(Nf_{op}(N)\log(N)))$ となります。
ここで、条件式 $f$ が満たしている性質を思い出します：
ある整数 $x > l$ に対し、$f(prod(l, x))=True$ であるとき、任意の整数 $l < y \leq x$ について $f(prod(l, y))=True$ である。また、$f(e) = True$ である。
これは $f$ が単調であることをいっています。すなわち、$True$ が $False$ に移り変わる境界は高々１箇所しかありません。このような場合には二分探索によって高速に答えを求めることができ、時間計算量は $O(f_{op}(N)(\log(N))^2)$ となります。
6.3. セグメント木上の二分探索とは
二分探索の解法では、$r$ を決め打ち $\rightarrow$ $prod(l, r)$ を取得 $\rightarrow$ $f(prod(l, r))$ を評価 $\rightarrow$ $r$ を決め打ち $\cdots$ というループをしました。
このとき、セグメント木はあくまでも $prod(l, r)$ を得るためのツールであり、ブラックボックス的に扱っています（セグメント木は、「$(l, r)$ を入力すると $prod(l, r)$ が出力されるもの」でしかありません。）。つまり、セグメント木による $prod(l, r)$ の取得と $r$ の探索は完全に分離したものになっています。
さて、今の状況をセグメント木の立場に立って考えてみましょう。するとこの状況は、「セグメント木の外で $r$ の二分探索を行なっている」と言えます。もっというと、セグメント木の外で何かしらの操作が行われ、それによって決まった「右端が様々な値をとる区間」が何度も入力されるという状況です。
ここでセグメント木の内部に目を向けてみると、「同じノードの参照や同じノード間の演算を何度も行う」という無駄が発生していることが分かります。
この無駄は、セグメント木が $prod(l, r)$ を取得するためのツールとなっていて、$r$ の探索と分離しているために起こっています。しかし、セグメント木は二分木であり、長さが２のべき乗の様々な区間の積を内包していることを考えると、セグメント木の内部で $r$ の二分探索ができる（＝セグメント木の内部のみで探索クエリが完結する）可能性が考えられます。
そして実際にこれは可能であり、セグメント木上の二分探索と呼ばれています。
セグメント木上の二分探索は先ほど見た無駄を無くしたものであり、時間計算量は $\boldsymbol{O(f_{op}(N)\log(N))}$ でセグメント木の外で二分探索を行う解法より $\log(N)$ が１つ少なくなります。
7. セグメント木上の二分探索の仕組み
では、セグメント木上の二分探索の仕組みを見ていきます。
7.1. 概要と通常の二分探索との違い
まず、セグメント木の外で二分探索を行う解法との違いを時間計算量の面から見ておきます。
前章で述べたように、セグメント木の外で二分探索を行う解法では時間計算量は $O(f_{op}(N)(\log(N))^2)$ でした。そしてその内訳は
- $prod(l, r)$ の取得 : $O(f_{op}(N)\log(N))$
- $r$ の探索 : $O(\log(N))$
です。$r$ の探索では探索範囲の初期値を $[l, N]$ とし、この範囲を半分にする操作を繰り返すので $O(\log(N))$ となっていました。
一方、セグメント木上の二分探索では時間計算量は $O(f_{op}(N)\log(N))$ であり、その内訳は
- $prod(l, r)$ の取得 : $O(f_{op}(N))$
- $r$ の探索 : $O(\log(N) + \log(N)) = O(\log(N))$
となります。$r$ の探索が $O(\log(N) + \log(N))$ となっているのが特徴的ですね。これは、ざっくりいうと
- 答えの大まかな位置を探す : $O(\log(N))$
- そこから徐々に細かく見ていき正しい答えを求める : $O(\log(N))$
という手順を踏んでいるからです。もう少し細かくいうと、$[l, r)$ の区間の長さが $1$ ($r = l + 1$) から始めて
- $f(prod(l, r)) = False$ となるまで $[l, r)$ の区間の長さを２倍し続ける
- 最後の２倍をする前の $r$ と今の $r$ の間に正しい答えがあるので（通常の）二分探索の要領で正しい答えを探す
というように探索を行います。そして、条件式 $\boldsymbol{f}$ が $\boldsymbol{True}$ となることが確定した区間についてはその積を保持しておくことで、$r$ を決めた際に $prod(l, r)$ を取得するためには１箇所のノードをみればよくなります。これによって $prod(l, r)$ の取得が $O(f_{op}(N))$ で出来るようになり、全体として $O(f_{op}(N)\log(N))$ となるわけです。
では、
STEP1 : 答えの大まかな位置を探す
STEP2 : 徐々に細かく見ていき正しい答えを求める
に分けて具体的な仕組みを見ていきます。
7.2. STEP1:答えの大まかな位置を探す
前節で述べたとおり、$f(prod(l, r)) = False$ となるまで $[l, r)$ の区間の長さを２倍し続けることで答えの大まかな位置を探します。
まず、$l$ が与えられたら葉に移動します(+size)。また、確定した区間の積を保持するための変数を用意しておきます。
そしてここから、$f(prod(l, r)) = False$ となるかもしくは区間の右端が $N$ となるまで区間を伸ばしていきます。このとき、実装する上での考え方は取得クエリと似ていて、
- 自身が左子ノードの時は親ノードが自身とより右のノードを共に含む
- 自身が右子ノードの時は親ノードが自身より左のノードを含む
なので自身が右子ノードの時のみ、そのノードを使用します。
擬似コードにすると次のようになります：
右探索クエリ（大まかな位置を求めるまで） (入力：探索の始点 l, 条件式 f){
l += size //葉に移動
sm = e //fがTrueとなることが確定した区間の積
while (True){
data[l]が右子ノード（もしくは根）になるまで{
親に移動
}
f(op(sm, data[l]))がTrueの場合{
sm = op(sm, data[l]) //確定した区間に加える
l += 1 //１つ右に移動
}
f(op(sm, data[l]))がFalseの場合{
終了(STEP2に続く)
}
f(prod(l, N))=Trueが確定した場合{
終了(max_rightはNで確定)
}
}
}
葉からスタートし各深さ（段）で一回の判定を行うので参照と演算の回数は $O(\log(N))$ 回です。
7.3. STEP2:徐々に細かく見ていき正しい答えを求める
STEP1での最後の２倍がやりすぎだということがわかっているので１.５倍はどうか？、1.25倍（もしくは１.７５倍）はどうか？...という具合に答えの範囲を絞っていきます。別の言い方をすれば、最後の１倍〜２倍の範囲で二分探索をします。
セグメント木の部分木もまた二分木になっているので、この二分探索は、今いるノードから葉ノードまで移動することに対応します。
具体的には、
- 左子ノードへ移動
- 今いるノードまで使って $f$ が $True$ なら１つ右に移動
という操作を繰り返します。上図を例にとれば、
- $3 \rightarrow 6$ に移動
- $f(17 \star 9 \star 5 \star 6) = True$ なので $7$ に移動
- $7 \rightarrow 14$ に移動
- $f(17 \star 9 \star 5 \star 6 \star 14) = False$ なので 移動しない
- $14 \rightarrow 28$ に移動
- $f(17 \star 9 \star 5 \star 6 \star 28) = True$ なので $29$ に移動
- $max\_right(l, f) = 29 - size = 13$
となります。
STEP１と合わせて擬似コードは次のようになります：
右探索クエリ(入力：探索の始点 l, 条件式 f){
l += size //葉に移動
sm = e //fがTrueとなることが確定した区間の積
while (True){
data[l]が右子ノード（もしくは根）になるまで{
親に移動
}
f(op(sm, data[l]))がFalseの場合{
//STEP2
葉ノードになるまで{
l <<= 1 //左子ノードに移動
f(op(sm, data[l]))がTrueの場合{
l += 1 //１つ右に移動
}
}
return l - size //終了
}
sm = op(sm, data[l]) //確定した区間に加える
l += 1 //１つ右に移動
f(prod(l, N))=Trueが確定した場合{
return N //終了(max_rightはNで確定)
}
}
}
STEP2ではあるノードから葉ノードまで移動するだけなのでノードの参照、演算の回数は $O(\log(N))$ 回です。
8. セグメント木上の二分探索の実装
では先ほどに擬似コードに沿って実装します。引数の条件式fはbool値を返す関数オブジェクトとして与えます。
STEP1において $f(prod(l, N))=True$ が確定するのは $l = (2のべき乗) - 1$ （図の$1, 3, 7, 15, 31$）でf(op(sm, self._d[l])) = True となった場合です。$f$ が $True$ のときには l が +1される（下記コードの下から３行目）ことを考えるとlが２のべき乗になった場合に$max\_right=N$ が確定することがわかります。
lが２のべき乗であることは「lの'1'の最下位ビットがlと等しい」ことと同値なのでこれで判断します。
例)
8ビットの2の補数表現で考えます。
① $l = 8$ のとき
\begin{aligned}
l \& (-l) &= (00001000)_2 \;\&\;(11111000)_2\\
&= (00001000)_2\\
&= 8
\end{aligned}
② $l = 9$ のとき
\begin{aligned}
l \& (-l) &= (00001001)_2 \;\&\; (11110111)_2\\
&= (00000001)_2\\
&= 1 \ne 9
\end{aligned}
def max_right(self, l, f):
assert 0 <= l <= self._n
assert f(self._e())
if l == self._n: return self._n
l += self._size # 葉に移動
sm = self._e() # 確定した区間の積を保持する変数
while True:
while l % 2 == 0: l >>= 1 # 右ノードになるまで
if not f(self._op(sm, self._d[l])):
# STEP2
while l < self._size:
l <<= 1
if f(self._op(sm, self._d[l])):
sm = self._op(sm, self._d[l])
l += 1
return l - self._size
sm = self._op(sm, self._d[l])
l += 1
if l & -l == l: break # f(prod(l, N))=Trueが確定
return self._n
左探索クエリを処理するメソッドmin_leftも同様に実装します。ただし、区間 $[l, r)$ が右開区間であるためSTEP1での１つ左に移動するタイミングやSTEP2での返り値がmax_rightと異なることに注意します。
def min_left(self, r, f):
assert 0 <= r <= self._n
assert f(self._e())
if r == 0: return 0
r += self._size
sm = self._e()
while True:
r -= 1
while r > 1 and r % 2: r >>= 1 # 左子ノードになるまで
if not f(self._op(self._d[r], sm)):
# STEP2
while r < self._size:
r = 2 * r + 1 # 右子ノードに移動
if f(self._op(self._d[r], sm)):
sm = self._op(self._d[r], sm)
r -= 1
return r + 1 - self._size
sm = self._op(self._d[r], sm)
if r & -r == r: break
return 0
9. 使用例
ここまで実装したものを使ってAtCoder Library Practice Contest J - Segment Treeを解いてみます。
問題文
長さ $N$ の整数列 $A = (A_1, A_2, \cdots, A_N)$ があります。
この数列に対して、$Q$ 個のクエリを処理してください。$i$ 番目のクエリでは、まず整数 $T_i$ が与えられます。 その後、$T_i$ に応じて以下の操作を行ってください。
$T_i=1$ : 整数 $X_i, V_i$ が与えられる。$A_{X_i}$ を $V_i$ で置き換える。
$T_i=2$ : 整数 $L_i, R_i$ が与えられる。$A_{L_i}, A_{L_i + 1}, \cdots , A_{R_i}$ の中での最大値を求める。
$T_i=3$ : 整数 $X_i, V_i$ が与えられる。$X_i \leq j \leq N, V_i \leq A_j$ を満たす最小の $j$ を求める。そのような $j$ が存在しない場合、$j = N + 1$ と答える。
制約（一部抜粋）
$0 \leq A_i \leq 10^9$
$0 \leq V_i \leq 10^9$
$T = 1$ は更新クエリ、$T = 2$ は取得クエリ、$T = 3$ は分かりにくいかもしれませんが右探索クエリになっているのでセグメント木を使って解けそうです。
まず、セグメント木が使えるか条件を確認します。
集合 $S$ はクエリ対象の整数列 $A$ が満たす制約から、「$0$ 以上 $10^9$ 以下の整数全体」とします。演算 $op$ は取得クエリから考えて $op(x, y) := max(x, y)$ とします。
1. 閉じているか
$max(x, y)$ は $x, y$ の大きい方を返すので明らかにこの演算は $S$ において閉じています。
2. 結合法則は満たすか
$x, y, z$ を $S$ の要素とし、$x \geq y, z$ の場合を考えます。
(i) $y \geq z$ のとき
\begin{aligned}
(x \star y) \star z &= max(max(x, y), z)\\
&= max(x, z)\\
&= x\\[4ex]
x \star (y \star z) &= max(x, max(y, z))\\
&= max(x, y)\\
&= x\\[4ex]
\therefore (x \star y) \star z &= x \star (y \star z)
\end{aligned}
(ii) $z \geq y$ のとき
\begin{aligned}
(x \star y) \star z &= max(max(x, y), z)\\
&= max(x, z)\\
&= x\\[4ex]
x \star (y \star z) &= max(x, max(y, z))\\
&= max(x, z)\\
&= x\\[4ex]
\therefore (x \star y) \star z &= x \star (y \star z)
\end{aligned}
$y, z$ が最大の場合も同様に示すことができます。したがってこの演算 $op$ は結合法則を満たします。
3. 単位元は存在するか
整数 $0$ は $S$ の要素であり任意の $x \in S$ に対して
max(0, x) = max(x, 0) = x
を満たすので単位元となります。が、後の事を考えて集合 $S$ に $-1$ を含めるようにしてこの $-1$ を単位元とします。
以上より
- 集合 $S$ : $-1$ 以上 $10^9$ 以下の整数全体
- 演算 $op$ : $op(x, y) := max(x, y)$
- 単位元 e : $-1$
とすると、$(S, op, e)$ はモノイドであり、セグメント木の使用条件を満たします。
セグメント木が使えることが分かったので、構築します。
def op(x, y):
return max(x, y)
def e():
return -1
n, q = map(int, input().split())
A = list(map(int, input().split()))
seg = SegTree(op, e, n, A)
各クエリを処理するコードを書きます。
t, a, b = map(int, input().split())
a -= 1
if t == 1: # 更新クエリ
seg.set(a, b)
elif t == 2: # 取得クエリ
print(seg.prod(a, b))
else: # 探索クエリ
...
$T = 1, 2$ はいいと思うので $T = 3$ を考えます。
$T = 3$ のクエリは
整数 $X_i, V_i$ が与えられる。$X_i \leq j \leq N, V_i \leq A_j$ を満たす最小の $j$ を求める。そのような $j$ が存在しない場合、$j = N + 1$ と答える。
でした。これが右探索クエリとなるように条件式 $f$ を適切に設定します。
この問題では $op(x, y) = max(x, y)$ なので $prod(l, r)$ は 区間 $[l, r)$ の要素の最大値です。また、
\begin{aligned}
&X_i \leq j \leq N, V_i \leq A_j を満たす最小の j\\[2ex]
\Leftrightarrow\; & (X_i \leq j \leq N, V_i > A_j を満たす最大の j) + 1
\end{aligned}
なので条件式を
f(x) := bool(x < V_i)
と定義することで右探索クエリに落とし込むことができます。
ここで、$f$ が満たすべき条件を満たしているか確認しておきます。
1. 単位元の入力に対して True を返すか
単位元は $-1$ であり、問題の制約から $V_i \geq 0$ が保証されているので
f(e) = True
です。ここが、単位元を $-1$ にした理由になります。
2. $\boldsymbol{f}$は単調か
$prod(l, r)$ は区間 $[l, r)$ の要素の最大値なので $r$ に対して単調非減少です。したがって、$f(x) = bool(x < V_i)$ は $True$ と $False$ が入れ替わる場所が高々１箇所、すなわち単調です。
以上より、定義した条件式 $f$ は右探索クエリの条件式として適当であることが確認できました。
あとはこれをコードにすれば解決です。
t, a, b = map(int, input().split())
a -= 1
if t == 1: # 更新クエリ
seg.set(a, b)
elif t == 2: # 取得クエリ
print(seg.prod(a, b))
else: # 探索クエリ
f = lambda x: x < b
print(seg.max_right(a, f) + 1)
これでこの問題を解くことが出来ました。
提出例
10. セグメント木上の二分探索の補足
10.1. 条件式が単調でない場合
条件式 $f$ が単調でない場合でも実装したコードは動作します。ただし、$max\_right(l, f)$ の返り値 $r$ は
f(prod(l, r)) = True かつ f(prod(l, r + 1)) = False
を満たす $r$ のうちのいずれか１つ（もしくは $r=N$）となります。（$min\_left$ も同様）
これは、通常の二分探索で解が複数ある場合を考えればわかると思います。
11. おわりに
今回はセグメント木をみてきました。
セグメント木はよく聞くデータ構造で応用の幅も広そうなので使いこなせるようになりたいですね。
説明の間違いやバグ、アドバイス等ありましたらお知らせください。