segtree

セグメントツリー 一点更新・区間クエリを高速で計算することが出来る。

初期化

G=segtree([0 for i in range(N)],func,ide_ele)

ここで、最初のリストは初期値である。ここは全部0である必要はない。必要に応じて変えてもいい。 例えば、Aだったり、list(range(N))だったりを入れる。 また、func,ide_eleは演算と単位元である。この演算はモノイドであることが要求される。 （注：モノイドとは、結合法則が成り立って、単位元が存在するような演算のことである。）

以下はセグ木に載せることのできる演算の例である。これはあくまで具体例なので他にもたくさんある。

セグ木関数	単位元	補足
add	0	足し算
times	1	掛け算
min	INF	最小値
max	-INF	最大値
gcd	0	最大公約数
lcm	1	最小公倍数
xor	0	排他的論理和
or	0	bitwise or
and	(2^N-1)	bitwise and（Nは制約に応じて十分大きな値を取る）
convolution	[1]	多項式の積（畳み込みを参照）
(a,b)*(c,d)->(ac,ad+b)	(1,0)	1次関数の合成,(a,b)はx->ax+bに対応
matrix	単位行列	行列の積

注：INFは、実装する際には適当な大きな数（ex:10の15乗）などを定める。ここで、INFの値の大きさが足りないと誤答の原因になるかもしれない。pythonだとmathにinfというものが実装されているが、バージョンが古いpythonでは使うことができない。

注2：足し算や掛け算は、掛け算＆modという組み合わせでも出てくるかもしれない。 modの世界で足し算や掛け算をしてもモノイドの性質は失われない。

注3：gcdを使う際にfrom math import gcdを使っているが、これもpythonのバージョンが古い場合エラーの原因となる。もしエラーが出た場合は

from fractions import gcd

で試してみるともしかしたらうまくいくかもしれない。

注4：1次関数の合成や、行列の積については、mod 998244353(or mod 10^9+7)上での演算として定める場合が多いと思われる。行列はk*kの積を計算するときにO(k^3)の計算量がかかるため、もし出るとしてもk=2,3あたりだろう。

実装上の注意

前節で例に挙げた各演算について、セグ木に載せる際には例えば以下のように書く。

#add
G=segtree(LIST,(lambda x,y:x+y),0)

#addの書き方その2
def add(x,y):
    return x+y
G=segtree(LIST,add,0)


#times
G=segtree(LIST,(lambda x,y:x*y),1)

#min
G=segtree(LIST,min,INF)

#max
G=segtree(LIST,max,-INF)

#gcd
from math import gcd
G=segtree(LIST,gcd,0)

#lcm
from math import gcd
def lcm(x,y):
    return (x*y)//gcd(x,y)
G=segtree(LIST,lcm,1)

# xor
G=segtree(LIST,(lambda x,y:x^y),0)

# or
G=segtree(LIST,(lambda x,y:x|y),0)

# and
N=30
G=segtree(LIST,(lambda x,y:x&y),(1<<N)-1)

関数はラムダ式で定義しても、defで定義してもどちらでも問題ない。

set

G.set(p,x)

p番目の値をxに変えることができる。

get

G.get(p)

p番目の値が返ってくるという関数である。

prod

G.prod(l,r)

[l,r)の範囲内での演算を求めた結果が返ってくる。 例えばセグ木関数がmaxだった場合max(A_l,...,A_{r-1})が返ってくる。 セグ木関数が足し算だった場合A_l+...+A_{r-1}が返ってくる。

関数の返り値は、このようなコードを実行したときの答えと同じである。

def prod(l,r):
    ans=ide_ele
    for i in range(l,r):
        ans=segfunc(ans,A[i])
    return ans

セグ木で書いた場合この区間クエリを高速に計算することができる。

all_prod

G.all_prod()

これはG.prod(0,N)と等価である。つまり、全区間での演算結果を求める。

max_right

G.max_right(l,f)

二分探索をする。ここで、始点はlであり、単調性のある関数fの実行結果が変わる切れ目を求める。

min_left

G.min_left(r,f)

二分探索をする。ここで、終点はrであり、単調性のある関数fの実行結果が変わる切れ目を求める。

update

内部処理用の関数なので使うことはあまりないかも…

debug

print(G)

2021年6月3日以降のバージョンでは、print関数を使うことでデバッグすることができます。 printをするとGの中身がlist形式として表示されます。

使用例

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

class segtree():
    #中略

N,Q=map(int,input().split())
A=[int(i) for i in input().split()]
G=segtree([i for i in A],max,-1)
for i in range(Q):
    t,a,b=map(int,input().split())
    if t==1:
        x,v=a,b;x-=1
        G.set(x,v)
    if t==2:
        l,r=a,b;l-=1;r-=1
        print(G.prod(l,r+1))
    if t==3:
        x,v=a,b;x-=1
        def f(t):
            if v>t:
                return True
            else:
                return False
        print(G.max_right(x,f)+1)