Pythonで遅延セグメントツリーの問題を解けるようにする！

## はじめに
Pythonメインで書かれた遅延セグメントツリーの記事が少ないので、書いてみることにしました。
この記事では、「PythonでAtCoderにある遅延セグメントツリーの基礎的な問題が解けるようになる」ことを目標にします。


事前知識として、次のことができる前提で進めます。
- 簡単なPythonのコードが読める。 (前半部分は概論チックでPythonを知らなくても読めます)
- 普通のセグメントツリーのお気持ち・雰囲気がなんとなく分かる。




## セグメントツリーの復習
セグメントツリーがどんなのだったか、いったん復習します。

区間の最大値を求めるセグメントツリーを用意しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/bba490b5-8f04-7427-2a1e-bb1b7c0b2a1a.png)

たとえば配列の $2$ ～ $6$ 番目の中の最大値を求めてくださいと言われたら...

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/778d9108-8662-8b68-cd09-7171dfb749a7.png)

図の赤の部分だけを見て、最大値を求めます。

また、配列の $4$ 番目の数字を $5$ に変えたいときは...
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/a8e7f7e5-9253-b7d3-6bf2-67217211c27b.png)

対応している葉の数字を変えて、そこから根までのノードを更新していきます。

ということで、まとめると、セグメントツリーの特徴は主に $2$ つです。

- 区間の演算が $O(logN)$ でできる。
- 値の $1$ 点更新が $O(logN)$ でできる。

では、次に遅延セグメントツリーのお話に入ります。

## 遅延セグメントツリーのできること
だいたい一緒です。遅延セグメントツリーのできることは主に $2$ つです。

- 区間の演算が $O(logN)$ でできる。
- 値の区間更新が $O(logN)$ でできる。

普通のセグメントツリーは、値の更新と言えば、配列のある $1$ 点だけを指して値を変更していました。
遅延セグメントツリーでは、「配列の $3$ 番目から $7$ 番目までをこんな感じで更新したい！」 みたいに、 $1$ 点だけを指さずに区間を指定して値を更新します。

区間の演算に関しては、普通のセグメントツリーと同じです。

## 遅延セグメントツリーのお気持ち
この問題を解きながら、遅延セグメントツリーについて見ていきましょう。
> 長さ $N$ の数列 $A=(A_1, A_2, \cdots ,A_N)$ が与えられます。
>  $Q$ 個のクエリを順に処理してください。クエリは以下の $2$ 種類です。
> - <code>1 l r v</code> : $i=l,l+1, \cdots ,r$ について、　$A_i \leftarrow A_i + v$ 
> - <code>2 l r</code> : $A_l, A_{l+1}, \cdots ,A_r$ の最大値を出力する。
>
> 制約1 : $1\leqq N,Q \leqq 2 \times 10^5$
> 制約2 : $1 \leqq A_i \leqq 10^9 (1 \leqq i \leqq N)$ 
> 制約3 : $1 \leqq v \leqq 10^9$
> 制約4 : $1 \leqq l,r \leqq N$


遅延セグメントツリーを用意しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/47c1496d-cb20-36c2-691f-5edf23ce3a5d.png)

「子が $2$ つあって、その最大値が自身の値になる」というところは、普通のセグメントツリーと全く同じです。
黄色の部分、lazyが何をするのかは、この先で説明していきます。

### 値の区間更新 (1)
<code>1 l r v</code> のクエリがとんできて、 $A_l, A_{l+1}, \cdots A_r$ それぞれについて $v$ だけ足すときに、遅延セグメントツリーではどのようなことをするのかを見ていきます。

普通のセグメントツリーの値更新は配列のある $1$ 点だけを更新するのに対して、遅延セグメントツリーでは配列のある区間の値を更新します。
たとえば、配列の $2$ ～ $4$ 番目の値に $21$ を足すとします。
すると、遅延セグメントツリーはこうなります。(赤丸の付いたノードが変わったところです)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/f9aa18b4-2367-fed3-da9e-1189687b917c.png)


配列の $2$ ～ $4$ 番目に対応するノードについて、dataを $+21$ して、さらにlazyも $+21$ しています。

lazyは、「本当ならこの下のノードたちにも $+21$ したいんだけど、全部やると時間がかかるから、 下のノードたちには $+21$ をせずに一旦ここに $+21$ の情報を保持しておく」みたいな役割です。

一部のノードのdataの値が変わったので、それに伴って上側のノードも変更しておきます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/73e51595-8655-1be4-f438-ebad4f4c79fc.png)


### 区間の演算 (1)
<code>2 l r</code> のクエリがとんできて、 $A_l, A_{l+1}, \cdots ,A_r$ の最大値を求めるとき、遅延セグメントツリーではどうするのかを見ていきます。

たとえば、$A_1, A_2, A_3$ の最大値を求めるとします。対応するのは矢印の付いたノードです。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/13951188-1d5a-f600-4c12-5124bb9b87f5.png)

しかし、値の更新をしていた際に、本来なら足されているはずの値を上側のノードのlazyに保持したままの可能性があります。実際に、葉の右側 $2$ つ(dataが $45,19$ のノード)には本来足されているはずの $21$ が足されておらず、上のノードのlazyに $21$ がたまっています。

なので、いきなり赤矢印のノードを見に行かずに、こういう手順で見ていきます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/f2ca42fe-ebc4-06c9-db90-fdfcc00fe2a0.png)


一番上の頂点(=根) から順番に、対応していたノード(赤矢印の付いていたノード)たちまで、lazyがたまってないかどうかチェックしていきます。
$1$ 番の矢印では、根のlazyの値は $0$ なので、その直下のノード $2$ つには何もしません。(lazyが $0$ なので、直下のノード $2$ つに $0$ を足したという認識でもokです)

これで赤矢印の付いたノードのうち左上の方の $data=33,lazy=0$ のノードには到達できました。
まだもう片方の赤矢印の付いたノードに到達できていないので、 $2$ 番の矢印です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/7a9e03d0-eab2-6bc9-9b59-59194899e8c2.png)

$data=66,lazy=21$ のノードから直下の $2$ つのノードに向けて、$lazy=21$ を反映させます。
具体的には、直下のノードのdataに $+21$ をして、直下のノードのlazyにも $+21$ をします。
今回は画像サイズの都合で $3$ 段の遅延セグメントツリーしか用意できてませんが、この下側にもまだノードがあるときを想像してもらえると、直下のノードのlazyも変更しないといけないのが分かるかと思います。
これで赤矢印の付いたノードすべてについて上側にあるlazyはすべて反映済みになり、赤矢印の付いたノードのdataはすべて正しい値になりました。

$A_1,A_2,A_3$ の最大値は、赤矢印の付いたノードのdataだけに注目して、 $\max(33,66)=66$ だと計算することができました。

### 値の区間更新 (2)
<code>1 l r v</code> のクエリが連続してたくさん飛んでくるときを考えます。
ゴチャついてきていたので、遅延セグメントツリーを元に戻しておきました。この図を基準に考えていきます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/501cfc68-9328-fe3f-9deb-eec1b7150dda.png)


まず、$l=1,r=2,v=35$ のクエリがとんできたとします。(赤丸の付いたノードが更新したノードです)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/c8ed63cb-de8d-54c0-4307-ed4e22329110.png)

上側のノードのdataも変更です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/1f5afba3-266d-0cd1-72b7-c8051a87e6e6.png)



次に、 $l=1,r=3,v=7$ のクエリがとんできたとします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/5644350f-9aaa-e492-1f55-cec591b9fc76.png)


対応するノードのdataに $+7$ すると同時に、lazyにも $+7$ します。
lazyに $0$ 以外の値が入っていても同じです。


上側のノードのdataも変更です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/e420a737-9904-ccde-d370-1ae8088e841c.png)

最後に、 $l=1,r=4,v=4$ のクエリがとんできたとします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/57a3d4c4-7576-8c78-c2db-5a5ea76a21ad.png)



### 区間の演算 (2)
たとえば、 $A_2$ の値が知りたいとします。
対応するノードは、次の図の赤矢印の付いたノードです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/c7422dd9-944b-c0f1-9af8-90bf164fc7cd.png)

ここには真の値は入っていなくて、上側にたまっているlazyをちゃんと反映させないといけないのでした。
ということで、一番上の頂点(=根)から順番に見ていきます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/6371ea7e-e551-d6ec-80bd-01d232a1d2e8.png)

根にあった $lazy=4$ を直下のノード $2$ つに反映させます。
dataを $+4$ するのはもちろん、lazyにも $+4$ します。これはlazyの値が $0$ でなくても同じです。
まだ赤矢印の付いたノードに到達できていないので、さらに下に行きます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/6b870225-493b-b86c-0073-a37409bebae4.png)

これで $A_2=58$ だと分かりました。

## Pythonで書いてみる
ここまで、遅延セグメントツリーのなんとなくの動き方をシミュレーションしてきました。
ここからは、それをもとに、Pythonで問題を解けるようにしていきます。

問題を解くのが目標なので、遅延セグメントツリー自体は他の人から拝借します。
今回は、ac-library-pythonにある遅延セグメントツリー[(これ)](https://github.com/not522/ac-library-python/blob/master/atcoder/lazysegtree.py)を使うことにします。 ac-library-pythonを自分のPCでつかえるようにする方法は、私の過去の記事 [(これ)](https://qiita.com/hyouchun/items/8a830952315666576c3d) に書いてあります。

ざっくりとしたつかい方は、こうです。
```Python
from atcoder.lazysegtree import LazySegTree

lazy_segtree = LazySegTree(op, e, mapping, composition, id_, v)
```

引数が $6$ つあります。これをうまく設定すると動いてくれます。
1つ1つ何を設定するのかを見ていきます。

### op: 区間の演算
これは普通のセグメントツリーと同じです。
区間の最大値を計算したいのであれば $\max$ を設定するし、最小公倍数lcmや排他的論理和xorも設定できます。
たとえば、排他的論理和xorを設定したいなら、こう書けばいいです。
```Python
def op(data1, data2):
    return data1 ^ data2
```

### e: opの単位元
単位元というとややこしいですが、こういう性質のものです。
$$op(data1, e) = data1$$
たとえば、opが最大値maxであれば、$e=-∞$　とすればいいですし、opが排他的論理和xorであれば、 $e=0$ とすればいいです。
演算しても全く意味のないものが $e$ です。

### mapping: 上のlazyを下のdataに作用させる

この状態の遅延セグメントツリーで、赤矢印の付いたノードを見たいとき、根から順番に辿ってlazyがないかチェックするのでした。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/c7422dd9-944b-c0f1-9af8-90bf164fc7cd.png)

根からその直下のノード $2$ つにlazyを作用させた結果がこうです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/6371ea7e-e551-d6ec-80bd-01d232a1d2e8.png)

根にあった $lazy=4$ を直下のノードのdataに作用させました。
この例では足し算するだけです。なので、mappingは次のように書きます。

```Python
def mapping(lazy_upper, data_lower):
    return data_lower + lazy_upper
```

### composition: 上のlazyを下のlazyに作用させる
先ほど貼り付けた遅延セグメントツリーの例では、上側のlazyの影響で下側のdataを変えただけではなく、下側のlazyも変わっていました。
この例では $2$ つのlazy同士を足し算するだけです。Pythonではこう書きます。

```Python
def composition(lazy_upper, lazy_lower):
    return lazy_upper + lazy_lower
```

### _id: mappingの恒等写像
恒等写像というとややこしいですが、こういう性質のものです。
$$mapping(data, id) = data$$
mapping関数のlazyの部分に入れて、そのままのdataが返ってくるものが_idです。

### v: 初期リスト
そのまま初期リストです。

## 問題を解いてみる (1)
(問題の再掲です)
> 長さ $N$ の数列 $A=(A_1, A_2, \cdots ,A_N)$ が与えられます。
>  $Q$ 個のクエリを順に処理してください。クエリは以下の $2$ 種類です。
> - <code>1 l r v</code> : $i=l,l+1, \cdots ,r$ について、　$A_i \to A_i + v$ 
> - <code>2 l r</code> : $A_l, A_{l+1}, \cdots ,A_r$ の最大値を出力する。
>
> 制約1 : $1\leqq N,Q \leqq 2 \times 10^5$
> 制約2 : $1 \leqq A_i \leqq 10^9 (1 \leqq i \leqq N)$ 
> 制約3 : $1 \leqq v \leqq 10^9$
> 制約4 : $1 \leqq l,r \leqq N$

ここまで書いてきたのをまとめると、こう書けます。

```Python
from atcoder.lazysegtree import LazySegTree


# 区間演算op
def op(data1, data2):
    return max(data1, data2)


# opの単位元 op(data1, e) = data1
e = -float("Inf")


# lazyをdataに作用させる
def mapping(lazy_upper, data_lower):
    return data_lower + lazy_upper


# lazyをlazyに作用させる
def composition(lazy_upper, lazy_lower):
    return lazy_upper + lazy_lower


# mapping(_id, data_lower) = data_lower
_id = 0

N = int(input())
A = list(map(int, input().split()))

lazy_segtree = LazySegTree(op, e, mapping, composition, _id, A)

Q = int(input())
for _ in range(Q):
    t, *q = map(int, input().split())
    match t:
        case 1:
            l, r, v = q
            lazy_segtree.apply(l - 1, r, v)
        case 2:
            l, r = q
            ans = lazy_segtree.prod(l - 1, r)
            print(ans)
```

## 問題を解いてみる (2)
典型90問の中にある、遅延セグメントツリーの問題を解いてみます。

https://atcoder.jp/contests/typical90/tasks/typical90_ac


### 問題要約
数直線上にレンガを積み上げていきます。
$i$ 番目のレンガは、 $l_i \leqq x \leqq r_i$ の範囲に積み上げます。
それぞれのレンガについて、高さを出力してください。

### シミュレーションする
入力例 $1$ はこのような図になります。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2469270/d0b90485-f092-287e-b9da-22b8f189daa3.png)


### 方針
$i$ 番目のレンガを積み上げるとき、 $l_i \leqq x \leqq r_i$ の区間の値の最大値を求めます。
その値に $+1$ したのが、 $i$ 番目のレンガの高さです。

### 引数を設定する
最初に図解していた遅延セグメントツリーでは、<code>1 l r v</code> のクエリで、配列の$l$ 番目から $r$ 番目まで $v$ を足していました。
今回は、指定した区間に $v$ を足すのではなく、指定した区間の値を一律に $v$ (=レンガの高さ)にしたいです。
なので、設定する引数も先ほどと少し変わってきます。

先に初期リストは、長さ $W$ で値がすべて $0$ の配列を用意すればいいです。この値はレンガの高さに対応しています。
```Python
# 初期リスト
lst = [0] * W
```

区間の演算では、指定した区間の値の最大値(=区間内のレンガの高さの最大値)を求めます。
```Python
# 区間演算op
def op(data1, data2):
    return max(data1, data2)

# opの単位元 op(data1, e) = data1
e = -1
```

lazyは、その区間全体に高さlazyのレンガが積まれていることを意味するとします。
何もなければ基本的には $lazy=-1$ で、たとえばあるノードのlazyの値が $3$ のとき、その下のノードのdataも $3$ にしないといけないけれど、いったんlazyにその情報をためているイメージです。

```Python
# lazyをdataに作用させる
def mapping(lazy_upper, data_lower):
    if lazy_upper == -1:
        return data_lower
    else:
        return lazy_upper


# lazyをlazyに作用させる
def composition(lazy_upper, lazy_lower):
    if lazy_upper == -1:
        return lazy_lower
    else:
        return lazy_upper


# mapping(_id, data_lower) = data_lower
_id = -1
```

これで引数 $6$ つをすべて設定できました。
まとめると、以下のコードになります。

<details> <summary>解答コード例 (長いので折りたたみ)</summary>

```Python
from atcoder.lazysegtree import LazySegTree

W, N = map(int, input().split())

# 初期リスト
lst = [0] * W


# 区間演算op
def op(data1, data2):
    return max(data1, data2)


# opの単位元 op(data1, e) = data1
e = -1


# lazyをdataに作用させる
def mapping(lazy_upper, data_lower):
    if lazy_upper == -1:
        return data_lower
    else:
        return lazy_upper


# lazyをlazyに作用させる
def composition(lazy_upper, lazy_lower):
    if lazy_upper == -1:
        return lazy_lower
    else:
        return lazy_upper


# mapping(_id, data_lower) = data_lower
_id = -1

lazy_segtree = LazySegTree(op, e, mapping, composition, lst)

for _ in range(N):
    l, r = map(int, input().split())
    l -= 1
    max_height = lazy_segtree.prod(l, r)
    lazy_segtree.apply(l, r, max_height + 1)
    print(max_height + 1)
```

</details>


## 問題を解いてみる (3)
最後に、AtCoder Library Practice Contest の中にある、遅延セグメントツリーの問題を解いてみます。

https://atcoder.jp/contests/practice2/tasks/practice2_k

### 問題要約
数列 $A=(a_0,a_1, \cdots , a_{N-1})$ が与えられます。$Q$ 個のクエリを順番に処理してください。クエリは $2$ 種類あって、以下の通りです。
<code>0 l r b c</code> : $i=l,l+1, \cdots r-1$ について、$a_i \leftarrow b \times a_i + c$
<code>1 l r</code> : $\sum_{i=l}^{r-1} a_i$ (mod $998244353$ ) を出力する。

### dataとlazyをどうするか
実は、今までのやり方だとうまくいきません。
天下り的になるんですが、dataには対応している区間の値の総和だけではなく、その区間の長さも必要です。要するに、(対応している区間の値の総和, 対応している区間の長さ) のタプルを dataとします。
dataやlazyは整数でもいいし、タプルでもいいし、なんなら自作のクラスでもいいわけです。

### 引数を設定する
ということで、LazySegTreeに入れる $6$ つの引数を設定していきます。

dataには、(対応する区間の値の総和, 対応する区間の長さ) のタプルが入ります。

```Python
# 区間の演算op
def op(data1, data2):
    sum1, length1 = data1
    sum2, length2 = data2
    next_sum = (sum1 + sum2) % MOD
    next_length = length1 + length2
    return (next_sum, next_length)

# opの単位元 op(data1, e) = data1
e = (0, 0)
```

lazyには $(b,c)$ のタプルを入れます。
たとえば、 $(a_0,a_1,a_2,a_3)$ の $4$ 要素それぞれに $a_i \leftarrow b \times a_i+c$ をすると、
$(a_0,a_1,a_2,a_3) \leftarrow (b \times a_0+c,b \times a_1+c,b \times a_2+c, b \times a_3+c)$ となり、その総和は $b \times (a_0+a_1+a_2+a_3) + c \times 4$ です。
$(a_0+a_1+a_2+a_3)$ の部分は、 $[0,3]$ の区間に対応するノードのdataの値です。
$c \times 4$ の $4$ の値は、 $[0,3]$ の区間の長さのことです。

これをPythonで書くと、次のようになります。

```Python
# lazyをdataに作用させる
def mapping(lazy_upper, data_lower):
    b, c = lazy_upper
    sum1, length1 = data_lower
    next_sum = (b * sum1 + c * length1) % MOD
    return (next_sum, length1)

# mapping(_id, data_lower) = data_lower
_id = (1, 0)
```

$a_i \leftarrow b_1 \times a_i+c_1$ の変換をした後、さらに $a_i \leftarrow b_2 \times a_i+c_2$ の変換をするときを考えます。
この $2$ つの操作をまとめると、 $a_i \leftarrow b_2 \times(b_1 \times a_i + c_1) + c_2 = b_1b_2 \times a_i + (b_2c_1+c_2)$ となります。 
これがcompositionにあたります。

```Python
# lazyをlazyに作用させる
def composition(lazy_upper, lazy_lower):
    b2, c2 = lazy_upper
    b1, c1 = lazy_lower
    next_b = b1 * b2 % MOD
    next_c = (b2 * c1 + c2) % MOD
    return (next_b, next_c)
```

最初のリストは、dataが(対応している区間の値の総和, 対応している区間の長さ) のタプルなので、次のように設定します。

```Python
A = list(map(int, input().split()))
lst = [(a, 1) for a in A]
```

これで遅延セグメントツリーをつくれました。
解答コードは以下の通りです。

<details> <summary>解答コード例 (長いので折りたたみ)</summary>

```Python
from atcoder.lazysegtree import LazySegTree

MOD = 998244353

N, Q = map(int, input().split())
A = list(map(int, input().split()))


# 区間の演算op
def op(data1, data2):
    sum1, length1 = data1
    sum2, length2 = data2
    next_sum = (sum1 + sum2) % MOD
    next_length = length1 + length2
    return (next_sum, next_length)


# opの単位元 op(data1, e) = data1
e = (0, 0)


# lazyをdataに作用させる
def mapping(lazy_upper, data_lower):
    b, c = lazy_upper
    sum1, length1 = data_lower
    next_sum = (b * sum1 + c * length1) % MOD
    return (next_sum, length1)


# mapping(_id, data_lower) = data_lower
_id = (1, 0)


# lazyをlazyに作用させる
def composition(lazy_upper, lazy_lower):
    b2, c2 = lazy_upper
    b1, c1 = lazy_lower
    next_b = b1 * b2 % MOD
    next_c = (b2 * c1 + c2) % MOD
    return (next_b, next_c)


lst = [(a, 1) for a in A]

lazy_segtree = LazySegTree(op, e, mapping, composition, _id, lst)

for _ in range(Q):
    t, *q = map(int, input().split())
    match t:
        case 0:
            l, r, b, c = q
            lazy_segtree.apply(l, r, (b, c))
        case 1:
            l, r = q
            ans = lazy_segtree.prod(l, r)[0]
            print(ans)
```

</details>


## 最後に
ここまで読んでくださってありがとうございました。

ここで取り扱った問題たちはどれも基礎的な問題で、遅延セグメントツリーの難しいところは「dataやlazyをどう設定するか」につきます(たぶん)。

あと、遅延セグメントツリーは定数倍がとても重くて、Pythonだと想定解通りでも実行制限時間に間に合わない悲しいケースもあります。なるべく計算が軽くなるような工夫をするようにしましょう。