Segment Tree Beatsの実装メモ (応用的な例題まわり)

前回の記事の続き
smijake3.hatenablog.com

コドフォのSegment Tree Beatsの説明の例題として挙がっているTask 3,4を実装したので、そのあたりの考え方含めて書いておく

A simple introduction to "Segment tree beats" - Codeforces

今回は考え方の説明がメインの記事。とりあえず自分が理解した(と思う)範囲で書いた。
例題については、方針の説明 + 実装コードになっている。

また、今回は主に以下を参考にして書いてある

区间最值操作与历史最值问题 - 道客巴巴 (3.2あたり)

※2019/05/28 計算量が間違ってる箇所があったため修正

セグ木の各ノードで区間情報を分けて管理する
- 区間chminクエリを扱う問題
- 区間chmin/chmaxクエリを扱う問題
値が変更されるたびに1加算する (Task 3)
- 実装コード
区間chminクエリを含む、2つの値の和の最大値を扱う (Task 4)

セグ木の各ノードで区間情報を分けて管理する

Segment Tree Beatsでは各ノードにおいて、区間が持つ情報を要素によって分けて管理しながらクエリを処理することができ、今回の例題も区間情報を分けて管理しながら処理をする。

ここでは、この考え方について説明する。

区間chminクエリを扱う問題

今回は、区間chminクエリとその他従来のクエリを扱う場合について考える。

また、以下の例題を元に説明していく。(この問題を解くためには区間が持つ情報を分割する必要がないが、説明のために分割することを考える)

長さ ${N}$ の数列 ${A}$ が与えられる。以下のクエリを合計 ${M}$ 回処理せよ。

$l \le i < r$ について $a_i$ の値を $\min(a_i, x)$ に更新

$l \le i < r$ について $a_i$ の値を $a_i + x$ に更新

$\displaystyle \sum_{l \le i < r} a_i$ を出力

制約: $N, M \le 10^5$

区間の最大値・非最大値による区間情報の分割

Segment Tree Beatsで区間chminクエリを処理する際、((厳密な)二番目の最大値) < x < (最大値) を満たす各ノードについて、その最大値の要素に新しい最大値 $x$ との差分を加算した上で区間が持つ情報(区間総和や区間最大値、等)を更新する。

この更新を行うノードでは 最大値の要素のみを(最大を維持するように)更新した上で、この最大値の要素に関する情報のみを更新することになる。逆に最大値以外の非最大値の要素や最大値に関係しない情報は更新しない。

この性質から、各ノードの区間に含まれる各要素 $a_i$ を最大値と非最大値の2種類に分け、区間が持つ情報(区間最大値や区間総和、遅延させている値(lazy tag)等)を最大値・非最大値ごとに分けて管理しながらクエリを処理できる。例えば、上の例題であれば総和を最大値の総和・非最大値の総和に分けて管理することができる。(この例では最大値の総和を $S_A$ 、非最大値の総和を $S_B$ とする)

イメージとしては下図のようになる。青枠には最大値の要素、緑枠には非最大値の要素が含まれ、この枠の単位で区間が持つ情報(総和や最大値、遅延させている値、等)を管理する。

f:id:smijake3:20190512034342p:plain — 要素を分割して管理するイメージ
(青枠が最大値、緑枠が非最大値。ノード上に枠が存在しない場合は要素が存在しない)

このように最大値と非最大値に分けて管理することにより、区間chminクエリは ((厳密な)二番目の最大値) < $x$ < (最大値) を満たすノードに対する、最大値のみに絞った区間加算操作として処理できる。

上の例題の区間chminクエリの更新について具体例で考えてみる。下の図は $a_i$ がとして8,8,4,6を含む区間を表し、この区間を $a_i \leftarrow \min(a_i, 7)$ で更新した例である。この場合、最大値(青枠)の各要素に区間加算クエリとして-1 (= 7 - 8)を加算し、この加算値を元に(最大値のみの)区間総和 $S_A$ を更新することができる。

f:id:smijake3:20190517232052p:plain — 区間情報を最大値と非最大値で分けて管理した上で、 $a_i \leftarrow \min(a_i, 7)$ を適用した例

各クエリで処理すること

このように分割すると各クエリに対して以下のように処理できる。

区間chminクエリ:
- 区間内の ((厳密な)二番目の最大値) < $x$ < (最大値) を満たすノードを更新
- 各ノードの最大値の要素に対し、新しい最大値 $x$ との差分を(遅延させながら)加算し、最大値に関する情報を(区間加算操作として)更新
区間加算クエリ(RAQ)や区間更新クエリ(RUQ)等の従来の操作クエリ
- 区間内の全てのノードを更新
- 各ノードの全ての要素を(遅延させながら)更新し、最大値と非最大値それぞれの情報を更新
区間総和クエリ(RSQ)や区間最大値クエリ(RMQ)等の従来の出力クエリ
- 区間を覆う各ノードの最大値に関する情報と非最大値に関する情報を合算した上で解を出力

この分割における区間chminクエリの処理イメージは以下の図のようになる。
((厳密な)二番目の最大値) < $x$ < (最大値) を満たすノードに対し、青枠(最大値)の要素に差分の加算処理を行った上で最大値に関する情報のみを(区間加算操作として)更新する。

f:id:smijake3:20190512131504p:plain — $\min(a_i, 7)$ で更新した例 (赤色が即更新、黄色が遅延させて更新)

$\min(a_i, 7)$ で更新した例 (赤色が即更新、黄色が遅延させて更新)

具体的な例として、上の例題では以下の処理を行うことになる。

区間chminクエリ:
- 区間内の ((厳密な)二番目の最大値) < $x$ < (最大値) を満たすノードを(遅延させながら)更新
- 各ノードの最大値の要素に対し、新しい最大値 $x$ との差分を加算し、最大値の総和 $S_A$ を加算値から更新
区間加算クエリ(RAQ):
- 区間内の全てのノードを(遅延させながら)更新
- 各ノードの全ての要素に $x$ を加算し、区間が持つ最大値の総和 $S_A$ 、非最大値の総和 $S_B$ をそれぞれ更新
区間総和クエリ(RSQ):
- 区間を覆う各ノードの、最大値と非最大値の総和 $S_A + S_B$ の和を求める

pushdownとupdateについて

情報を最大値・非最大値に分けて管理する場合、pushdownやupdateの時に区間の最大値・非最大値の関係を考慮しながら伝搬する必要があり、この伝搬関係は左右の子ノードの最大値の値によって決まる。

updateでは、親ノードの情報は左右の子ノードの最大値・非最大値の情報(合わせて4種類の情報)を元に更新する。
具体的に、最大値に関する情報は左右の子ノードのうち、より大きい最大値を持つ子ノードの最大値の情報から更新する。この時、両方同じ最大値を持つ場合は両方の最大値の情報をマージして更新する。非最大値に関する情報は、それ以外の情報をマージして更新する。

f:id:smijake3:20190515014512p:plain — updateのイメージ
(左図は "左子ノードの最大値 > 右子ノードの最大値"、右図は "左子ノードの最大値 = 右子ノードの最大値")

例えば、上の例題のように最大値の総和 $S_A$ と非最大値の総和 $S_B$ を管理している時に、子ノードの各総和から親ノードの各総和をupdateする場合を考える。
この時、左右の子ノードの最大値によって以下のように場合分けしてupdateする。

左子ノードの最大値 > 右子ノードの最大値
- 親ノードの $S_A$ = 左子ノードの $S_A$
- 親ノードの $S_B$ = 左子ノードの $S_B$ + 右子ノードの $(S_A + S_B)$
左子ノードの最大値 < 右子ノードの最大値
- 親ノードの $S_A$ = 右子ノードの $S_A$
- 親ノードの $S_B$ = 左子ノードの $(S_A + S_B)$ + 右子ノードの $S_B$
左子ノードの最大値 = 右子ノードの最大値
- 親ノードの $S_A$ = 左子ノードの $S_A$ + 右子ノードの $S_A$
- 親ノードの $S_B$ = 左子ノードの $S_B$ + 右子ノードの $S_B$

pushdownにおける遅延値(lazy tag)の伝搬はupdateの逆方向に行う。
最大値に関する遅延値は、左右の子ノードのうち親ノードと等しい最大値を持つ子ノードへ(最大値が等しい場合は両方へ)伝搬させる。それ以外には、非最大値に関する遅延値を伝搬させる。

f:id:smijake3:20190515014956p:plain — pushdownのイメージ
(左図は "左子ノードの最大値 > 右子ノードの最大値"、右図は "左子ノードの最大値 = 右子ノードの最大値")

最大値に関する遅延値として $a_i$ の最大値の差分のみを子ノードに伝搬すればよい場合は、pushdownで $a_i$ の最大値を更新する際の差分の値で伝搬できるため、最大値の差分加算を遅延させるための別のlazy tagを持って処理する必要はない。

例題の実装

上の例題の実装を一応載せておく。
基本的に区間の情報を分割する前に比べて、総和を分割しただけの実装になる。

計算量は $O(N \log N + M \log^2 N)$

実装コード(Gist): Segment Tree Beatsの例題実装 (Task 3, Task 4まわり) · GitHub

区間chmin/chmaxクエリを扱う問題

区間chmin/chmaxクエリ両方を扱う問題でも同様に区間の情報を分けて解くことができ、最大値・最小値・それ以外の3種類に分けることができる。

この時注意すべきなのは区間の要素が (最大値) = (最小値) のケースで、特別に最大値と最小値に同じ要素が含まれると考える。このケースの実装では、要素を重複して解に含めないように注意する。(総和で要素を2回足す、など)

f:id:smijake3:20190516010452p:plain — 最大値・最小値・それ以外の3つに分けて管理するイメージ

updateやpushdownも区間chminクエリのみの場合と同じ考え方で処理する

update:
- 親ノードの最大値(最小値)の情報は、左右の子ノードの中での最大値(最小値)の情報から更新する
- それ以外の要素の情報は、左右の子ノードの中での最大値・最小値以外の情報をマージして更新する
pushdown:
- 親ノードの最大値(最小値)に関する遅延値は、左右の子ノードで等しい最大値(最小値)を持つ子ノードへ伝搬する
- それ以外には、非最大値かつ非最小値の要素に関する遅延値を伝搬する

f:id:smijake3:20190518113641p:plain — update, pushdownでの更新・伝搬関係のイメージ

値が変更されるたびに1加算する (Task 3)

長さ ${N}$ の数列 ${A}$ が与えられる。また、最初全て0の長さ $N$ の数列 ${B}$ が与えられる。
以下のクエリを合計 ${M}$ 回処理せよ。

$l \le i < r$ について $a_i$ の値を $\max(a_i, x)$ に更新

$l \le i < r$ について $a_i$ の値を $\min(a_i, x)$ に更新

$l \le i < r$ について $a_i$ の値を $a_i + x$ に更新

$\displaystyle \sum_{l \le i < r} b_i$ の値を出力

(各クエリ後、そのクエリで $a_i$ の値が変わった全ての $i$ について、 $b_i$ に1を加算)

制約: $N, M \le 10^5$

数列 $A$ をクエリで処理するたびに数列 $B$ の値が増加する問題。
この問題では $b_i$ の区間総和に答えるために、各ノードでは $b_i$ の区間総和を管理する。

区間加算クエリ(RAQ)のみであれば、 $x \neq 0$ の場合に区間内の $b_i$ に1を一様に加算すればよい。しかし区間chminクエリは $a_i > x$ を満たす全ての $i$ について $b_i$ に1を加算する必要がある。

ここで各ノードにおいて以下のように $b_i$ を3種類に分け、それぞれの総和を管理することを考える。

$a_i$ が区間最大値の $b_i$ (この総和を $S_A$ とする)
$a_i$ が区間最小値の $b_i$ (この総和を $S_B$ とする)
$a_i$ が上記2つ以外の $b_i$ (この総和を $S_C$ とする)

すると区間chmin(chmax)クエリにおける $b_i$ の更新処理は、各ノードの $a_i$ の最大値(最小値)が変化した時に、最大値(最小値)に関する全ての $b_i$ に(遅延させながら)1を一様加算する処理になり、その分を総和 $S_A$ ( $S_B$ )に加算すればよくなる。

f:id:smijake3:20190516013857p:plain — この問題におけるデータ管理イメージ
(枠内は $b_i$ の値を示し、この枠単位の総和 $S_A,S_B,S_C$ を管理する)

この問題におけるデータ管理イメージ
(枠内は $b_i$ の値を示し、この枠単位の総和 $S_A,S_B,S_C$ を管理する)

f:id:smijake3:20190517001624p:plain — 上の例に対し min(a_i, 7) を取った例 (赤色が即更新、黄色が遅延させて更新)
この更新に合わせて $S_A$ も更新する

上の例に対し min(a_i, 7) を取った例 (赤色が即更新、黄色が遅延させて更新)
この更新に合わせて $S_A$ も更新する

実装的には、この問題では総和を分けて管理する必要性はないため、各ノードでは全ての総和の和 $S_A + S_B + S_C$ の値を管理すればよい。

また、区間加算クエリは区間に含まれる全ての $b_i$ に1を一様加算するだけになる。
そのため、遅延させながら一様加算するための遅延値として

$a_i$ が最大値となる $b_i$ に加算する用 (区間chminクエリで加算)
$a_i$ が最小値となる $b_i$ に加算する用 (区間chmaxクエリで加算)
全ての $b_i$ に加算する用 (区間加算クエリで加算)

の3種類を管理すればよい。

計算量は $O(N \log N + M \log^2 N)$

実装コード

全体の実装(Gist): Segment Tree Beatsの例題実装 (Task 3, Task 4まわり) · GitHub

区間chminクエリを含む、2つの値の和の最大値を扱う (Task 4)

２つの数列 $A$ と $B$ の各 $i$ における和 $a_i + b_i$ の区間最大値を扱う問題。
Task 4では両方の数列で区間chminクエリを扱うが、まずは簡単な例として片方のみが区間chminクエリを扱う例を考えてみる。

1つの数列のみで区間chminクエリを処理する問題

長さ ${N}$ の数列 ${A}$ と数列 ${B}$ が与えられる。以下のクエリを合計 ${M}$ 回処理せよ。

$l \le i < r$ について $a_i$ の値を $\min(a_i, x)$ に更新

$l \le i < r$ について $b_i$ の値を $x$ に更新

$\displaystyle \max_{l \le i < r} a_i + b_i$ を出力

制約: $N, M \le 10^5$

この問題では $a_i + b_i$ の最大値が必要となるため、この情報を各ノードで維持する必要がある。

もし数列 $A$ が従来の区間操作(RAQ, RUQ)のみであれば、 $a_i + b_i$ に関して区間最大値のみを持つことで処理できる。
しかし区間chminクエリでは $a_i > x$ となる $i$ だけ $a_i + b_i$ の値を更新するため、(遅延させて更新する関係上) $a_i + b_i$ に関して区間最大値のみ持つだけでは処理できない。例えば、 $a_i$ が小さくなることで $a_i$ が更新されなかった $a_i + b_i$ が新たな最大値になる場合があり、これを考慮するのは難しい。

f:id:smijake3:20190518020501p:plain — $a_i+b_i$ の最大値のみを管理した時に更新に失敗するケース
(区間の $a_i$ を $\min(3, a_i)$ で更新した際、 $a_i + b_i$ の区間最大値を8 (=2+6)に更新できない)

そこで $a_i$ が非最大値の時の $a_i + b_i$ の区間最大値も考慮できるように、 $a_i$ が最大値の時と非最大値の時で区間情報を分けて管理することを考える。
具体的に、各ノードで以下のペア情報を管理する。(ペアではなく $a_i + b_i$ の値をそのまま持っても処理できるが、今回はペアで持って処理する)

$a_i$ が区間最大となる $i$ の中で $a_i + b_i$ が最大となるペア $P_A = (a_i, b_i)$
$a_i$ が区間非最大となる $i$ の中で $a_i + b_i$ が最大となるペア $P_B = (a_i, b_i)$

以下の図は数列 $A = (1,1,3,2)$ 、数列 $B = (8,2,3,6)$ における各ノードの管理イメージである。"×"の箇所は該当するペアが存在しないことを意味しており、実装では $(-\infty, -\infty)$ を保持させる。

f:id:smijake3:20190512114505p:plain — 各ノードにおける管理イメージ
(各ノードの左が $P_A$ 、右が $P_B$ )

各ノードにおける管理イメージ
(各ノードの左が $P_A$ 、右が $P_B$ )

そして以下のように処理する。

の区間chminクエリ:
- 更新対象となるノードにおける、 $P_A$ に含まれる $a_i$ を $x$ に更新
の区間更新クエリ(RUQ):
- 更新区間に含まれる全てのノードについて
  - $P_A$ に含まれる $b_i$ を $x$ に更新
  - $P_B$ の $a_i$ を二番目の最大値、 $b_i$ を $x$ に更新
の区間最大値クエリ:
- 区間を覆う各ノードについて $P_A$ と $P_B$ から $a_i + b_i$ の最大値を求め、全てのノードにおける最大値を出力

計算量は $O((N+M) \log N)$

実装コード

全体の実装(Gist): Segment Tree Beatsの例題実装 (Task 3, Task 4まわり) · GitHub

区間chminクエリが処理される2つの列を処理する問題

次に本題であるTask4を考えてみる。

長さ ${N}$ の数列 ${A}$ と数列 ${B}$ が与えられる。以下のクエリを合計 ${M}$ 回処理せよ。

${l \le i < r}$ について ${a_i}$ の値を ${\min(a_i, x)}$ に更新

${l \le i < r}$ について ${b_i}$ の値を ${\min(b_i, x)}$ に更新

${l \le i < r}$ について ${a_i}$ の値を ${a_i + x}$ に更新

${l \le i < r}$ について ${b_i}$ の値を ${b_i + x}$ に更新

$\max_{l \le i < r} a_i + b_i$ を出力

制約: ${N, M \le 3 \times 10^5}$

一つ前の例のように $a_i + b_i$ の区間最大値を扱えるようにするために、各数列 $A$ 、 $B$ がそれぞれ最大値・非最大値の場合に分けて管理することで解くことができる。

具体的に $A$ と $B$ の最大値・二番目の最大値に加えて、各ノードでは以下の4種類のペア $(a_i, b_i)$ の情報を管理する。(これも同様に $a_i + b_i$ の値をそのまま持っても処理できる)

$a_i$ と $b_i$ が最大値となる $i$ の中で $a_i + b_i$ が最大値となるペア $P_{11}$
$a_i$ が最大値、 $b_i$ が非最大値となる $i$ の中で $a_i + b_i$ が最大値となるペア $P_{10}$
$a_i$ が非最大値、 $b_i$ が最大値となる $i$ の中で $a_i + b_i$ が最大値となるペア $P_{01}$
$a_i$ と $b_i$ が非最大値となる $i$ の中で $a_i + b_i$ が最大値となるペア $P_{00}$

以下の図は数列 $A = (1, 2, 3, 2)$ 、数列 $B = (5, 3, 5, 1)$ を $a_i + b_i$ が最大となる4種類のペアを管理するイメージである。
("×"の箇所は該当するペアが存在しないことを示し、値としては $(-\infty, -\infty)$ を保持させる)

f:id:smijake3:20190512023639p:plain — 4種類のペアの管理イメージ
(各ノードの左上が $P_{11}$ 、右上が $P_{10}$ 、左下が $P_{01}$ 、右下が $P_{00}$ )

4種類のペアの管理イメージ
(各ノードの左上が $P_{11}$ 、右上が $P_{10}$ 、左下が $P_{01}$ 、右下が $P_{00}$ )

f:id:smijake3:20190512024220p:plain — ノードで管理する最大値ペアの更新イメージ
(赤枠: $a_i$ に対する区間chminクエリで更新するペア, 青枠: $b_i$ に対する区間chminクエリで更新するペア)

ノードで管理する最大値ペアの更新イメージ
(赤枠: $a_i$ に対する区間chminクエリで更新するペア, 青枠: $b_i$ に対する区間chminクエリで更新するペア)

数列 $A$ に対する区間chminクエリでは $a_i$ が最大値となる $P_{11}$ と $P_{10}$ の $a_i$ の値を更新し(上図の赤枠)、数列 $B$ に対する区間chminクエリでは $b_i$ が最大値となる $P_{11}$ と $P_{01}$ の $b_i$ を更新する(上図の青枠)。
数列 $A$ と数列 $B$ に対する区間加算クエリでは、全てのペアの $a_i$ 、 $b_i$ の値を更新すればよい。

そして、 $a_i + b_i$ の最大値を出力する際は、区間内の4種類のペアの中の $a_i + b_i$ の最大値を出力すればよい。

計算量は $O(N \log N + M \log^2 N)$

実装コード

全体の実装: Segment Tree Beatsの例題実装 (Task 3, Task 4まわり) · GitHub

区間chminクエリが処理されるK個の列を処理する場合

この問題を拡張し、各数列で区間chminクエリを扱う $K$ 個の数列における、同じ $i$ の総和の区間最大値を処理する場合について考えることができる。
この場合、各数列について最大値・非最大値からなる全ての組み合わせの最大値ペア( $2^K$ 個)を管理する必要があるため、計算量は $O(2^K (N \log N +M \log^2 N))$ となる。