lazysegtree - shakayami/ACL-for-python GitHub Wiki

lazysegtree

区間一括更新ができるセグ木。実行時間を食うので必要がない場合はセグ木を使ったほうがいい。

初期化関数

初期化をするためには3つの関数を予め定義することが要求される。 その3つとは、

  • operate
  • mapping
  • composition

である。これが何かについては、そもそもの遅延セグ木の使用法の話につながる。

遅延セグ木に乗る元は集合Gの要素である。ここで、Gは普通のセグメントツリー同様にモノイドであることが要求される。さらに、Fという写像全体の集合も出てくる。ここで、FはMap(G,G)(つまり、GからGへの(準同型とは限らない)写像)の部分集合となっていて、これは

  • 恒等写像の存在
  • 写像の合成に関して閉じている

ことが要求される。

このとき、上記の遅延セグ木三点セットに対応する関数は以下を意味している。

operate

operateはG×G→Gというモノイドにおける積の写像を意味している。

たとえば、(a_l * ... * a_{m-1})と(a_m * ... * a_{r-1})という2つの元を入力とした時、 (a_l * ... * a_{m-1} * a_m * ... * a_{r-1})という元を出力とするような関数である。

しかし、普通のセグメント木と比べて単純なものではなく、後述のmappingがうまく定義できるようにしなくてはいけない。詳細はmappingにて

mapping

これは、Fの元であるfとGの元であるxに対して、

mapping(f,x)=f(x)

を返す関数である。 つまり、mappingはF×G→Gという関数である。 ここで注意しなければいけないのは、xは区間での情報を持っていることである。 xはa_l * ... * a_{r-1}という形になっている場合、 mapping(f,x)はf(a_l)* ... * f(a_{r-1})という出力になっている。 つまり、mappingで演算をうまく定義するために、余分に情報を添加する場合がある。これはRange Affine Range Sum ( https://atcoder.jp/contests/practice2/tasks/practice2_k )がわかりやすい例である。

composition

これは、Fの元であるf,gを取ってきたときに、

composition(f,g)=f○g

という関数である。これは写像の合成を意味している。 これは、F×F→Fという関数である。

初期化

seg=lazysegtree(LIST,operate,e,mapping,composition,identity)

ここで、

  • operate,mapping,compositionは前述の通り
  • eはGにおける単位元
  • identityはFにおける単位元(つまり恒等写像) を意味している。

set

一点更新

seg.set(p,x)

p番目の要素をxに変える

get

一点取得

seg.get(p)

p番目の値を取得する

prod

区間取得

seg.prod(l,r)

区間[l,r)内での演算結果を返す。

all_prod

全区間取得

seg.all_prod()

区間[0,N)内(つまり全体)での演算結果を返す。

apply

区間更新

seg.apply(l,r,f)

区間[l,r)内の要素a_l,...,a_{r-1}を、f(a_l),...,f(a_{r-1})に変更する。

max_right,min_left

二分探索をする。普通のセグメントの場合と同じ。

使用例

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

class lazy_segtree():
    # (中略)
N,Q=map(int,input().split())
a=[int(i) for i in input().split()]
ans=[]
mod=998244353
def operate(a,b):
    return ((a[0]+b[0])%mod,a[1]+b[1])
def mapping(f,x):
    return ((f[0]*x[0]+x[1]*f[1])%mod,x[1])
def composition(f,g):
    return ((f[0]*g[0])%mod,(g[1]*f[0]+f[1])%mod)

seg=lazy_segtree([(i,1) for i in a],operate,(0,0),mapping,composition,(1,0))
for i in range(Q):
    seq=tuple(map(int,input().split()))
    if seq[0]==0:
        dummy,l,r,b,c=seq
        seg.apply(l,r,(b,c))
    else:
        dummy,l,r=seq
        ans.append(seg.prod(l,r)[0])
for line in ans:
    print(line)

この例においては、Gはmod 998244353全体の集合(つまり位数998244353の有限体)となっている。 演算は足し算である。

また、Fはx->bx+cという写像全体の集合である。 ここで注意しなければいけないのは、xがa_l,...,a_{r-1}という区間を持っていて、fがx->bx+cという関数のとき、

mapping(f,x)=(ba_l+c)+...+(ba_{r-1}+c)=b(a_l+...+a_{r-1})+c(r-l)

というようになる。ここで、mappingを計算するときに区間の長さの情報を持ってないと計算することが出来ない。 よってなし崩し的にもoperateにも影響が出て、x+yという演算では不十分であり、 (演算結果,区間の長さ)という情報をセグ木上で持ってないといけないのである。

さて、このままだとTLEになってしまうのである。解決方法としては、セグ木上で情報を持つ際に、タプルではなく数値にすることである。INFで下駄を履かせて、a*INF+bという形で数値を持てばa,bという数値を持つことが出来る。たとえば実装例は以下のようになる。

class lazy_segtree():
    # (中略)

import sys
input=sys.stdin.readline

N,Q=map(int,input().split())
a=[int(i) for i in input().split()]
ans=[]
mod=998244353
def operate(a,b):
    a0,a1=a>>32,a%(1<<32)
    b0,b1=b>>32,b%(1<<32)
    return (((a0+b0)%mod)<<32)+a1+b1
def mapping(f,x):
    f0,f1=f>>32,f%(1<<32)
    x0,x1=x>>32,x%(1<<32)
    return (((f0*x0+x1*f1)%mod)<<32)+x1
def composition(f,g):
    f0,f1=f>>32,f%(1<<32)
    g0,g1=g>>32,g%(1<<32)
    return (((f0*g0)%mod)<<32)+((g1*f0+f1)%mod)

seg=lazy_segtree([(i<<32)+1 for i in a],operate,0,mapping,composition,1<<32)
for i in range(Q):
    seq=tuple(map(int,input().split()))
    if seq[0]==0:
        dummy,l,r,b,c=seq
        seg.apply(l,r,(b<<32)+c)
    else:
        dummy,l,r=seq
        ans.append(seg.prod(l,r)>>32)
for line in ans:
    print(line)

遅延セグ木の使用例をもう一つ紹介する。 (https://atcoder.jp/contests/practice2/tasks/practice2_l)

先ほど紹介した情報を踏まえて、セグ木上にどの情報を乗せればいいのかを考えてみよう。正解は以下の通り。

class lazy_segtree():
    # (中略)
import sys
input=sys.stdin.readline
 
N,Q=map(int,input().split())
a=[(0,1,0) if i=="0" else (0,0,1) for i in input().split()]
def op(x,y):
    return (x[0]+y[0]+x[2]*y[1],x[1]+y[1],x[2]+y[2])
def mapping(f,x):
    if f==0:
        return x
    else:
        return (x[1]*x[2]-x[0],x[2],x[1])
def composition(f,g):
    return f^g
seg=lazy_segtree(a,op,(0,0,0),mapping,composition,0)
for i in range(Q):
    t,l,r=map(int,input().split())
    l-=1
    if t==1:
        seg.apply(l,r,1)
    else:
        ans=seg.prod(l,r)
        print(ans[0])

正解は

(区間内での転倒数,区間内の0の個数,区間内の1の個数)

である。区間をマージする際にどのような情報があれば積を計算できるかを考えることがコツである。