AtCoder Library を読んでアルゴリズムを勉強：セグメント木（一括更新版）

<details><summary>前回と同じ前書き</summary><div>

競技プログラミング [AtCoder](https://atcoder.jp/home) では、コンテストでよく使うアルゴリズムをライブラリにした **AtCoder Library (ACL)** を[公開していて](https://atcoder.jp/posts/517)、特に C++ ならコンテスト内で[利用できます](https://atcoder.jp/posts/535)。

私は Ruby を使っているため直接は利用できず、ライブラリを模写する必要がありました。その過程で学んだアルゴリズムを整理してまとめます。

</div></details>

---

今回は**セグメント木**の続きで、区間内の要素の一括更新が可能なほうを扱います。

- AtCoder Library の `lazysegtree`
	- [ソースコード](https://github.com/atcoder/ac-library/blob/v1.3/atcoder/lazysegtree.hpp)
	- [ドキュメント](https://github.com/atcoder/ac-library/blob/v1.3/document_ja/lazysegtree.md)

名前に lazy が付いている通り**遅延評価**（ lazy evaluation ）を用いて実現しています。遅延評価するために、基本版に無かった数学的な条件やアルゴリズムの手順などが増えているので、それらを中心にまとめます。

※ 基本版と被る説明は省略しているので、[前回の記事](6aaf0c9157cf96a64e59)を適宜参照してください。

## TL;DR
<!--
セグメント木で区間の一括更新ができるものについて、 AtCoder Library （ACL）の実装を元に勉強・整理しました。
-->

- 区間和の計算・1要素の更新だけでなく、区間の更新も O(log n) で処理できます。
- 区間内の各要素を更新するための関数には、**「区間和にも直接適用できること」**という条件があります。
	- 簡単な例は分配法則を満たすもの（整数の和に対する積など）です。
	- 分配法則を直接は満たさないときでも、要素を構造体などにして付加情報を持たせると対応できる場合があります。
- データ構造では、セグメント木のノードを遅延評価できるように情報を持たせます。
	- 子ノードに未適用の関数（あるいはそのパラメータ）を親ノードに記憶させます。
	- メソッド内ではノードを参照する前に、記憶した関数を最上段のノードから順に反映する必要があります。

## 利用場面
ここで言う「区間の一括更新」というのは、ある区間の要素に任意の値を代入するのではなく、**各要素に関数を適用して値を置き換える**ものです。例えば「区間の各要素に100を加算する」などです。

<details><summary>例：路線バスの予想所要時間・改</summary><div>

> 前回と同じく路線バスの予想所要時間を求めたいとします。データとして各バス停間の時間が与えられていて、道路状況に合わせてデータを更新します。セグメント木を使えば「任意区間の時間の積算」と「データの更新（1箇所）」の両方をバランス良い速さで行えます。
>
> ところで、データの更新で「渋滞の影響を受ける連続した区間を、全て同じ時間だけ加算したい」という要望が出てきたとしましょう。データの更新をするには、各バス停間についてひとつずつ処理することになり、該当区間が長いと大変です。一方で所要時間を求める観点では、**各バス停間は興味なく**、必要なら `加算する時間 * 加算される区間の長さ` と補正するだけで済みます。何か更新処理を短縮する方法は無いでしょうか？

</div></details>

基本版のセグメント木では、元データの長さが n のとき1要素の更新にかかる時間は O(log n) でした。長さ m の区間について愚直に更新すれば O(m log n) かかります。あるいは更新したデータからセグメント木を作り直せば O(n) です。
以下の図は n = 8 で `a[2...8]` を更新する場合を示しています。セグメント木の最下段を更新してから親ノードも再計算し、結果的にほとんどのノードで更新が発生しています（青い要素のうち太線で囲ったもの）。

![segment-tree-lazy-update0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/4ec0b971-809d-2e2d-223a-a4f08ae5b2d4.png)

一括更新版では更新方法に条件を付けることで、**区間の長さ m に関わらず O(log n) で更新できます**。基本版にあった区間和の計算なども、何倍か遅くなりはするものの O(log n) のままです。

### 扱える計算
基本版と同じく、モノイドと呼ばれる代数的構造（型・二項演算 `op` ・単位元 `e` の適切な組）なら扱えます。

- `op(x, y)` の戻り値の型は `x`, `y` と同じであること
- `op(op(x, y), z) == op(x, op(y, z))` であること
- 次式を満たす `e` が存在すること： `op(e, x) == x && op(x, e) == x`

ただし区間の一括更新に対応するには、型は単純な数値でなく構造体などで付加情報を持つように修正しなければいけないことがあります。

### 更新用の関数の性質
要素に適用する関数は何でもいいわけではなく、いくつか条件が課されています。これらの条件の意味を見ていきます。（条件を課す理由は[後の節](#更新方法の条件とデータ構造)でまとめます）

ACL のドキュメントでは f 自体を関数としていますが、実装やサンプルを読む際はこれを**関数 `F` とパラメータ `f` に分解**した方が良さそうです。例えば「100を加算する」なら、加算が関数、100がパラメータ、という具合です。

※ ACL の `F` は `f` の型ですが、本記事の `F` は ACL の `mapping` と同じ意味で使っています。

- `x_ = F(f, x)`
	- 要素を更新する関数 `F` （ ACL の `mapping` ）は、引数にパラメータ `f` と元の要素 `x` をとり、新しい要素 `x_` を返します
	- `F` の戻り値の型は当然 `x` の型と同じですが、パラメータの型は別でも構いません
- `op(F(f, x1), F(f, x2)) == F(f, op(x1, x2))`
	- **関数 `F` は、要素だけでなく区間和も直接更新できます**[^op-n]
	- この際にパラメータ `f` は全て同じです
	- 小ネタとして、 `F` は要素の単位元を更新しません： `F(f, e) == e` <br>（実装が間違っていないかの軽い検証に使えます）
- `F(f1, F(f2, x)) == F(G(f1, f2), x)`
	- 同じ要素を連続して更新する場合、代わりにパラメータを合成しておき一発で処理できます
	- この合成を関数 `G` （ ACL の `composition` ）で定義します
		- `G` は結合法則を必ず満たします： `G(G(f1, f2), f3) == G(f1, G(f2, f3))`
		- `G` は交換法則を満たすとは限りません<br>（更新クエリの順序を変えると結果が変わりうるということです）
- `F(H(), x) == x`
	- 要素を更新しない働きのパラメータが存在します<br>（これを関数 `H` （ ACL の `id` ）に生成させることとします）
	- パラメータの単位元でもあります： `G(H(), f) == f && G(f, H()) == f`

2番目の性質を満たす `F` を見つけることが重要（かつ厄介）であり、 `G`, `H` は割と自然に作れます。

[^op-n]: `op` を可変長引数に拡張して考えると、「区間和を直接更新できる」という意味がわかりやすいです。

#### 具体例
2番目の性質を満たす `F` の簡単な例は、 **`op` と[分配法則](https://ja.wikipedia.org/wiki/分配法則)が成り立つもの**です。

- `op` が加算、 `F` が乗算<br> → `(f * x1) + (f * x2) == f * (x1 + x2)`
	- 各要素を定数倍してから合計するのと、合計してから定数倍するのは、同じ結果です
	- 要素をベクトルや行列、パラメータを正方行列に拡張したりもできます（線型代数）
- `op` が最小値、 `F` が加算<br> → `min(f + x1, f + x2) == f + min(x1, x2)`

---

一方で、 **`op` が加算のときに `F` も加算にしたい**ことはよくありますが、素直に考えると2番目の性質を満たしません： `(f + x1) + (f + x2) != f + (x1 + x2)` 。

この場合、要素を構造体などに変えて**区間和に追加情報を持たせる**と解決します。具体的には、区間和がカバーしている要素数の情報があれば、 `F` はその数だけ一度に加算できるようになります。

```javascript
// 要素や区間和の型を作り直し（区間の要素数 size を追加）
class Item {
	constructor(value, size) {
		this.value = value;
		this.size = size;
	}
}

// 型が変わったので、二項演算（と単位元）も作り直し
function op(x1, x2) {
	return new Item(x1.value + x2.value, x1.size + x2.size);
}

// 更新用の関数をうまく定義できる
function F(f, x) {
	return new Item(x.value + f * x.size, x.size);
}
```

---

さらに別の例で、 `op` は加算で **`F` は指定値で上書き**というものを考えてみます。区間和を直接更新する場合は、先ほどと同じ考え方で要素数を考慮した値で上書きすればいいです。また、パラメータの合成 `G` は単に新しい方を採用するだけです（交換法則が成り立たない例にもなっています）。

```javascript
// 要素の型は先ほどと同じ `Item`
function F(f, x) {
	return new Item(f * x.size, x.size);
}

function G(f1, f2) {
	return f1;
}
```

しかしこれらの関数は、「要素を更新しないパラメータ」が無く `H` を定義できません。 `F` は `x.value` の情報を消してしまいますし、 `G` は `f2` を消してしまいます。

このようなときの簡単な対処法は**更新しないための特別なパラメータを付け加える**ことです。ズルに感じますが、条件を満たしさえすれば関数の中で条件分岐していても問題ありません。
パラメータは他と区別できれば何でもいいので、今回は `null` を追加してみます。コード変更は早期リターンの追加（とnullの許可）だけで、定型的には以下のようになります。

```javascript
function F_(f, x) {
	if (f == null) return x;
	return new Item(f * x.size, x.size);
}

function G_(f1, f2) {
	if (f1 == null) return f2;
	if (f2 == null) return f1;
	return f1;
}

function H_() {
	return null;
}
```

#### 補足：関数オブジェクトを使った理解
<details><summary>（高階関数が苦手でない人向けです）</summary><div>

関数をオブジェクトにできる言語なら、**パラメータ `f` 自体を関数にする**方法も考えられ、 ACL のドキュメントの f と同じ意味になります。この場合、 `F`, `G`, `H` は引数や戻り値で関数を扱うことになるので「高階関数」となります。

```javascript:高階関数（アロー関数を使った場合）
// ACL のドキュメントに沿った定義
const F = (f, x) => ( f(x) ); // mapping
const G = (f1, f2) => ( x => f1(f2(x)) ); // composition
const H = () => ( x => x ); // id

// 単位元（恒等関数）を特殊値にして、 G による関数呼び出しのネスト軽減
const F_ = (f, x) => {
	if (f == null) return x;
	return f(x);
};
const G_ = (f1, f2) => {
	if (f1 == null) return f2;
	if (f2 == null) return f1;
	return x => f1(f2(x));
};
const H_ = () => ( null );
```

これらの高階関数は任意の関数 `f` を扱えるので、様々な問題で使い回せます。一方で「区間和を直接更新できる」という条件は `F` が担保しないため、それぞれの `f` 自身が満たす必要があります。また `G` は関数呼び出しをネストするだけで最適化しないため、 `f` が関数でない場合と比べて処理が遅くなると考えられます。

</div></details>

## アルゴリズム
※ 具体的な計算イメージが欲しいときは、以下の設定が簡単でわかりやすいです。

```javascript
// モノイド：整数の加算
function op(x1, x2) { return x1 + x2; } // 何個も繋げると区間和
function e() { return 0; } // 要素の単位元

// 一括更新用の関数：整数倍
function F(f, x) { return f * x; }
function G(f1, f2) { return f1 * f2; }
function H() { return 1; } // パラメータの単位元
```

### （基本版の）区間和の計算から考察
例えば n = 8 で `a[2...8]` の6個の和を求めたいとき、セグメント木では親ノードの値を使って `op(s[5], s[3])` を計算すれば済みます。前回の記事と同様に図示すれば以下のようになります。

![segment-tree-lazy-sum0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/68b9f23f-35b5-64c8-90cb-a5cd6479739b.png)

こう見ると**区間和の計算時は、最下段 `s[8...16]` の値は必ずしも必要ではなく、しかも参照するノード数が少なく済む**ことがわかります。ならば一括更新する際も、実は**後で使うノードだけを更新すればいい**のではないでしょうか。

簡単な例は、一括更新する区間と、*その直後に*和を求める区間が同じ場合です。 `a[2...8]` の例であれば、それらをカバーする `s[5]` と `s[3]` を直接更新して子は放置します。親ノードの再計算は少なく済みそうなのでやっておきます。

![segment-tree-lazy-update1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/7f3333d3-832d-b6d4-f8d5-1ca5f0906b4f.png)

この方法は一括更新を短縮できて良さそうですが、子を放置したということはその部分のデータは使い物にならなくなります。別の区間の参照や更新もできるためには、放置した部分も何とかして利用可能にする必要があります。そこで遅延評価の考え方を使い、必要になったときに**後からでも利用可能にする**仕組みを作ります。

### 更新方法の条件とデータ構造
まず大前提として最下段（元データのコピー）を個別に更新することはしません。要素毎に異なる処理（例えば新しい配列データで上書き）を用意すると、セグメント木を横方向に走査することになり、最悪 O(n) かかってしまいます。それを避けるには**要素毎の処理を同じにする**ことになり、ひとつの関数を各要素に適用する方法が落としどころです。<br>→ `x_ = F(f, x)`

前節の検討を踏まえると、最下段だけでなく**親ノードに対して直接更新をかけたい**です。これを成り立たせるためには `(更新した要素)の和 == (要素の和)の更新` という条件が必要です。<br>→ `op(F(f, x1), F(f, x2)) == F(f, op(x1, x2))`

親ノードを直接更新した場合、後から子ノードを更新する方法が必要になりますが、このときも同じ条件が役立ちます。**親ノードの更新に使ったパラメータと同じもので子ノードを更新すればいい**だけです。そのためにはパラメータをノードに記憶させる必要があるため、パラメータも同じ添字を持つ木構造として管理します（最下段は子ノードが無いため不要です）。データ構造に `f[]` を加えた図を以下に示します（元の配列 `a[]` は省略します）。

![segment-tree-lazy-design.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/d6eff802-36f8-1904-852f-d0f00f28fd21.png)

---

各ノードに「子ノードを更新するためのパラメータ」を記憶させましたが、複数の更新が重なった場合への対処が必要です。わかりやすい方法は、記憶領域をキューにしてパラメータを貯めていき、更新を適用するときに*全て*取り出して順次処理する方法です。欠点として更新が蓄積すると、メモリを多く消費する上に `F` を何度も実行することになってしまいます。
これを改良した方法としてパラメータの合成が考えられます。**順次処理した場合と同じ結果を得られるパラメータ**を作れるなら、キューでなくひとつ記憶させるだけで済みます。<br>→ `F(f1, F(f2, x)) == F(G(f1, f2), x)` [^f1-f2]

キューで記憶する代わりにパラメータを合成するとなれば、キューが空のときに相当する「更新しない」パラメータが必要です。<br>→ `F(H(), x) == x`

セグメント木の初期状態では全ノードで子ノードの更新が不要なので、パラメータの木 `f[]` には全て `H()` を入れておきます。

[^f1-f2]: `G` の引数のどちらが先に適用するパラメータかはきちんと決めておく必要があります。 ACL では第2引数が先で、そこに第1引数を合成します。数学の[合成関数](https://ja.wikipedia.org/wiki/写像の合成)でも、だいたいこの順番で書きますが逆にする記法もあります。

### 要素の一括更新
この処理はもちろん基本版に無かった機能であると共に、一括更新版の戦略が詰まったものになっています。そのためこれだけ理解できれば他のメソッドについても理解できます。

処理は3段階になっています。

0. 利用するノードの値を評価（計算）します
0. 区間をカバーできるノードを選んで関数を適用します
0. 親ノードの値を再計算します

「利用するノード」は後の処理を考えて初めて決定できるので、 2 → 3 → 1 の順に見ていきます。

#### ノードへの関数適用
ここまでで何度も見た通り、「関数 `F` は区間和を直接更新できる」という条件を使って、親ノードに対して関数を直接適用します。区間をカバーする親ノードを見つけるには、区間和の計算時と同じ走査方法を使えばいいです。

更新しなかった子ノードを後から利用可能にするために、「子ノードに反映していないパラメータ」を同じノードに記録しておきます。ただし既にパラメータが記録されている可能性があるので、上書きするのではなく関数 `G` で合成します。

`a[2...8]` をパラメータ `f0` で更新する場合の動作を以下に図示します。

<!--
![segment-tree-lazy-apply.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/f88a230b-fcff-0424-d98b-055a8d324a19.png)
-->
![segment-tree-lazy-apply.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/196c46de-c187-f7bd-71e9-75c53d5851cd.png)

#### 親ノードの再計算
セグメント木は親ノードが子ノードの結果を持っておく構造です。そのため子ノードを更新したら親ノードを再計算する必要があります。

単純に考えれば、更新を適用したノードそれぞれについて親ノードを再計算します。いま扱っている例では `s[5]` と `s[3]` を更新したので、以下の計算をすればいいです。

- `s[5]` → `s[2] = op(s[4], s[5]); s[1] = op(s[2], s[3])`
- `s[3]` → `s[1] = op(s[2], s[3])`

![segment-tree-lazy-update.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/8f756c00-add8-d19c-fd87-b043a1432bd4.png)

よく見てみると `s[1]` で**再計算が重複**しています。何度計算しても結果は同じままなので安全ですが、明らかに無駄な計算は省きたいです。

実は**関数を適用した端のノードだけを起点に更新すれば全て反映できる**性質があります。前節の図で、区間左端を示している "\[" は `s[5]` と `s[3]` を通過し、右端を示している "\]" はまっすぐ上へ行きました。この場合は `s[5]` だけで十分ということになります。

#### 利用するノードの値を評価
関数の適用や親ノードの再計算のためには、当然ノードの値を参照する必要があります。しかし遅延評価の戦略をとっているということは、**ノードに関数が反映されていない可能性があり**すぐ参照してはいけないということです。

参照するためには遅延評価を実行します。**信頼できるのは最上段だけ**なので、そこに記録したパラメータを子ノードに反映していきます。ノード `k` から子ノードに反映するメソッドを `push(k)` とすれば、 `s[5]` と `s[3]` は以下の順番に計算します。

- `s[5]` → `push(1); push(2)`
- `s[3]` → `push(1)`

子ノードに反映したら親ノードに記録しておいたパラメータ `f[k]` は用済みになります。2回以上実行しても結果が変わらない（冪等になる）ように、 `f[k] = H()` とリセットしておく必要があります。

この操作を図に示します。最初は添字 `1` だけ参照可能で、他は未反映です（点線のノード）。 `push` を実行することで反映していけます。 `s[5]` と `s[3]` を参照可能にするよう反映していくと、前節までと同じ状態になり関数適用・再計算ができます。

![segment-tree-lazy-push.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/857667/c2f75a6e-aef5-b77c-1435-dc9c94cdad6b.png)

ところで、これを前節の再計算と比較すると、**走査方向が上下逆になっているだけで全く同じ流れ**です。添字 `1` での計算が重複している点まで同じです。そのためこちらも関数を適用する端のノードだけを起点に更新すれば十分です。

### 基本版にあった操作
他の操作のアルゴリズムは基本版と同じです。異なるのは要素の更新が反映されていない点であり、**値を参照する前に必ず反映**しなければなりません。

大半の操作はどのノードを参照するかすぐ特定できるため、**基本版に前処理を追加しただけ**の作りになっています。前処理の細かいバリエーションとして、 `get` や `set` は最下段を利用するため全ての段で計算し、一方で区間和 `prod` は最下段まで利用するとは限らないのでif文が付いています。

ただし二分探索（ `min_left` と `max_right` ）は事情が異なります。探索するにつれて参照すべきノードが判明するので、前処理だけで済ませることはできません。それでも無駄を減らすために工夫されています。

- 1段階目（区間の拡大）では関数の適用と同じ論理が使え、区間の端（引数で与えた側）だけを準備すれば、1段階目に参照しうる全てのノードが利用可能になります。
- 2段階目（絞り込み）では木を1段ずつ下るため、直上のノードが既に利用可能です。そのため最上段から再評価する必要はありません。

基本版を理解した状態なら、このあたりに気をつければ ACL の実装を読み解けるはずです。