セグメント木も遅延セグメント木も、ある一定の区間における演算結果が見たい時に有用なのは共通しているが、
- セグメント木は個々の値を更新したいとき
- 遅延セグメント木はある一定の区間においていっぺんに値を変更する
点で異なる。実装例を交えながらより詳しく見ていく。
セグメント木
セグメント木は、配列構造を二分探索木のような形状に変更することで、ある特定の値が更新されたときに特定の範囲を演算する回数を普通の配列よりも減らす役割を担っている。
例えば[1,2,3,4]という配列が与えられており、値を更新しながら与えられた範囲の最大値を得たい場合は
- 初期化時に各区間の最大値をあらかじめ計算しておく
- (1,2)の最大値 -> 2
- (3,4)の最大値 -> 4
- 子ノードの計算結果より(2,4)の最大値 -> 4
- 値が更新された場合
- 2番目の値が2から5に更新されたとする
- (1,5)の最大値 -> 5
- (3,4)の最大値は既に計算されているので(4,5) -> 5
- 1番目から3番目までの最大値を計算したい場合
- (1,5)の最大値は既に計算されている
- （
atcoder.segtreeでは便宜上ここで(3,-INF)の最大値が計算されている） - 以上の計算結果より(5,3)の最大値 -> 5
セグメント木は、配列中のある特定の値を更新しつつ、総和や最小値、最大値などのある特定の期間の演算結果が見たい場合に使う。
問題例
長さ$N$の数列$a_0,a_1,...,a_{N-1}$に$Q$個のクエリが飛んできます。処理してください。
0 p x: $a_p \leftarrow a_p + x$1 l r: $\sum^{r-1}_{i=l} a_i$ を出力する。
- $A[i] \leftarrow x$
- $\sum^{r}_{i=l} A[i]$を計算
# pip install git+https://github.com/not522/ac-library-python
from atcoder.segtree import SegTree
N,Q = list(map(int,input().split()))
A = list(map(int,input().split()))
st = SegTree(op=lambda a,b:a+b, e=0, v=A)
# op: 2つの値が入ってきたときの演算
# e: 単位元(op(a,e) = aとなるようなe)
# v: 配列
for _ in range(Q):
query = list(map(int,input().split()))
match query[0]:
case 0:
p,x = query[1:]
# 配列のp番目の値をupdate
pre_x = st.get(p)
st.set(p,pre_x+x)
case 1:
l,r = query[1:]
# l~r-1の値で演算を行う
print(st.prod(l,r))
遅延セグメント木
遅延セグメント木は、セグメント木を拡張させて特定の区間更新を少ない計算量でできるようにした木構造である。
セグメント木での例と同じで[1,2,3,4]が与えられて最大値を計算する場合を考える
-
配列の初期化
- データ構造の初期化はセグメント木と同じ
- lazy構造には後述するマッピングにおける単位元(
mapping(e,data)=dataとなるようなe)をあらかじめ入れておく、この場合は$-INF$
-
1番目~3番目の値に10が代入される場合
- 1番目と2番目の親ノードと3番目のlazyに10を入れる。
- 1番目と2番目の値を親ノードのlazyに応じて変換する。これが
mappingの役割である。今回の場合は子ノードの値とlazyのうちどちらか大きい方を子ノードのデータとする。 - 1番目と2番目のlazyを親ノードのlazyに応じて変換(
composition)したら親ノードのlazyは0に戻す。この場合は子ノードのlazyと親ノードのlazyのうちどちらか大きい方を子ノードのlazyとする。 - （3番目のノードについても仮想の親ノードを作って1番目、2番目と同様の処理を行うことで値がちゃんと10に更新される）
- あとは該当区間の親(祖先)ノードの演算結果を普通のセグメント木と同様に更新していく。
-
与えられた区間の演算はセグメント木と同じ
- この演算をする際にlazyにまだ残っているデータがあったらその都度
mappingやcomposition関数を呼び出して必要な範囲の値を計算する。
- この演算をする際にlazyにまだ残っているデータがあったらその都度
問題例1
$W$個の正方形のマスが左右に並んだ水平な部分があります。最初、すべての部分について、高さは$0$です。ここに高さ$1$の直方体のレンガを$N$個、順番に積みます。高さ$h$の面に接着したレンガの上面の高さは$h+1$になります。
i番目に積むレンガは、左から$L_i$番目から$R_i$番目のマスをちょうど覆うように置きます。このとき、レンガが覆う範囲の中で最も高い水平な面で接着します。
各レンガについて、上面の高さを求めてください。
from atcoder.lazysegtree import LazySegTree
W,N = list(map(int,input().split()))
V = [0 for _ in range(W)]
def op(data1,data2):
return max(data1,data2)
def mapping(lazy_upper, data_lower):
return max(lazy_upper, data_lower)
def composition(lazy_upper, lazy_lower):
return max(lazy_upper,lazy_lower)
e = -float("INF")
id_ = -float("INF")
lst = LazySegTree(op,e,mapping,composition,id_,V)
for _ in range(N):
l,r = list(map(int,input().split()))
maximum_range = lst.prod(l-1,r) # 配列でいうところの[l-1:r]と同義
lst.apply(l-1,r,maximum_range+1)
print(maximum_range+1)
問題例2
長さ$N$の整数列$a_0,a_1,...,a_{N-1}$が与えられ。$Q$個のクエリが飛んできます。処理してください。
0 l r b c: 各$i=l,l+1,...,r-1$について、$a_i \leftarrow b \times a_i + c$1 l r: $\sum_{i=l}^{r-1} a_i mod 998244353$を出力する。
この例ではデータ構造やlazyが単なる値ではなくtupleになっている場合について扱っている。LazySegTreeではデータやlazyが必ずしもint型である必要はない。
from atcoder.lazysegtree import LazySegTree
N, Q = list(map(int,input().split()))
A = list(map(int,input().split()))
A = [(a,1) for a in A]
mod = 998244353
def op(data1,data2):
sum_a1,l1 = data1
sum_a2,l2 = data2
#print(f"op:{data1},{data2} -> {((a+aa)%mod, l+ll)}")
return ((sum_a1+sum_a2)%mod, l1+l2)
def mapping(up_lazy,data):
sum_a, l = data
b, c = up_lazy
#print(f"mapping:{up_lazy},{data} -> {((b * sum_a + l*c) % mod,l)}")
return ((b * sum_a + c*l) % mod,l)
def composition(up_lazy,lazy):
b1,c1 = lazy
b2,c2 = up_lazy
#print(f"composition:{up_lazy},{lazy} -> {(bb * b, bb * c + cc)}")
return ((b2 * b1)%mod, ((b2 * c1)%mod + c2)%mod)
lst = LazySegTree(op,(0,0),mapping,composition,(1,0),A)
for _ in range(Q):
query = list(map(int,input().split()))
match query[0]:
case 0:
l,r,b,c = query[1:]
lst.apply(l,r,(b,c))
case 1:
l,r = query[1:]
print(lst.prod(l,r)[0])
#print("A:",[lst.get(i) for i in range(N)])
なぜわざわざtuple型を扱って初期値を(a,1)のような形状にしているのかというと、mapping関数に秘密がある。
ではまず(a,1)ではなくaとして格納されていたとしよう。
この場合、初期化時に各親ノードのデータ構造に子ノードの総和が格納された状態になる。
したがって、mapping関数が呼ばれた際に親ノードのlazyを子ノードの要素に作用させるとき、子ノードが末端のノードでない限りはデータ構造は複数の要素の総和が格納されている状態になるので
\sum \hat{a}_i = \sum (b \times a_i + c) = b\sum a_i + L \times c \\
が本来格納されなければならない。($L$は該当ノードにまとめられている要素の数)
しかし、データ構造にaしか格納されていないと$L$の値が分からないので正しく情報が更新されない。そのため初期化値を(a,1)として親ノードの値を計算するときに$L$の値を足し算するとmapping関数が正しく作用する状態になる。
ちなみに今回のケースにおけるcompositionは、万が一子ノードのlazyが消化されないまま親ノードのlazyの情報が降りてきたとき、
\begin{align}
b_{new} \times \sum \hat{a_i} + c_{new} = b_{new} \sum (b_{old} \times a_i + c_{old}) + c_{new} = (b_{new} + b_{old})\sum a_i + (b_{new} \times c_{old} + c_{new}) = B \sum a_i + C
\end{align}
となるように更新されなければならないので、
\begin{align}
B &= b_{new} + b_{old}\\
C &= b_{new} \times c_{old} + c_{new}
\end{align}
となる。
問題例3
長さ$N$の数列$A=(A_1,A_2,...,A_N),B=(B_1,B_2,...,B_N)$が与えられます。
$Q$個のクエリが与えられるので、順に処理してください。
クエリは次の3種類です。
1 l r x: $A_l, A_{l+1}, ...,A_r$ に $x$を加える。2 l r x: $B_l, B_{l+1}, ...,B_r$ に $x$を加える。3 l r: $\sum_{i=l}^{r} (A_i \times B_i)$を998244353で割った余りを出力する。
- $X_i$に$(A_i,B_i,A_i \times B_i, 1)$を格納
- $X_l,...,X_{r}$に対して値を更新
- $X_l,...,X_{r}$に対して演算処理
from atcoder.lazysegtree import LazySegTree
# ※できるだけ余り演算を行う回数を減らす & あまり大きい値を保持しないようにする
N,Q = list(map(int,input().split()))
A = list(map(int,input().split()))
B = list(map(int,input().split()))
AB = [(a , b, a*b, 1) for a,b in zip(A,B)]
mod = 998244353
# 2つの子ノードを用いて親ノードの値を演算
def op(data1,data2):
sum_a = (data1[0] + data2[0]) % mod
sum_b = (data1[1] + data2[1]) % mod
sum_ab = (data1[2] + data2[2]) % mod
length = data1[3] + data2[3]
return (sum_a, sum_b, sum_ab, length)
# op(data1,e) = data1
e = (0,0,0,0)
# Aに足すのなら(x,0), Bに足すのなら(0,x)をlazyに保持するとする
def mapping(lazy_upper, data_lower):
# lazy_upperとdata_lowerは既に余り演算を終えているものだと仮定
to_add_a, to_add_b = lazy_upper
sum_a, sum_b, sum_ab, length = data_lower
# 親ノードのlazyから子ノードのdataを更新
sum_ab = (sum_ab + (to_add_b * sum_a) % mod + (to_add_a * sum_b) % mod + (length * to_add_a * to_add_b) % mod) % mod
sum_a = (sum_a + length * to_add_a) % mod
sum_b = (sum_b + length * to_add_b) % mod
print(f"{lazy_upper},{data_lower} -> {(sum_a, sum_b, sum_ab, length)}")
return (sum_a, sum_b, sum_ab, length)
def composition(lazy_upper, lazy_lower):
# 親ノードのlazyから子ノードのlazyを更新（更新された後lazy_upperはid_となる）
to_add_a, to_add_b = lazy_upper
to_add_aa, to_add_bb = lazy_lower
print(f"{lazy_upper},{lazy_lower} -> {((to_add_a + to_add_aa) % mod, (to_add_b + to_add_bb) % mod)}")
return ((to_add_a + to_add_aa) % mod, (to_add_b + to_add_bb) % mod)
# mapping(id_, data_lower) = data_lower
id_ = (0,0)
lst = LazySegTree(op,e,mapping,composition,id_,AB)
for _ in range(Q):
t, *q = map(int,input().split())
match t:
case 1:
l,r,x = q
# l-1~rのlazyに(x%mod,0)を入れる
lst.apply(l-1,r,(x%mod,0))
case 2:
l,r,x = q
# l-1~rのlazyに(0,x%mod)を入れる
lst.apply(l-1,r,(0,x%mod))
case 3:
l,r = q
#A×Bの総和を計算
ans = lst.prod(l-1,r)
print(ans[2])
"""
print("===A===")
print([lst.get(i)[0] for i in range(N)])
print("===B===")
print([lst.get(i)[1] for i in range(N)])
print("===AB===")
print([lst.get(i)[2] for i in range(N)])
"""
セグメント木と遅延セグメント木との違い
-
特定の値更新か、広い範囲における値更新か
セグメント木は個々の値更新しかできないので、広い範囲で共通した値更新を行う場合は遅延セグメント木を用いた方がよい。 -
lazyを使うか使わないか
遅延セグメント木にはdata(大元の配列の要素)だけではなく、値更新の際にlazyという部分も使う。値更新処理が行われると該当範囲のlazy部分に値(上のコードにおける(x%mod,0)や(0,x%mod))が入力され、その値に応じてmappingやcompositionを用いて子ノードのdataやlazy部分が更新される。