はじめに
AtCoderをやっているとセグ木貼って終わりとか強い人がTwitterでつぶやいているのを度々見かけるのですが、自前で持っていなかったのでこの際実装して理解、応用につなげようと言う記事です。
今コンテスト中でたまたま検索でたどり着いて今すぐ実装がほしいひとは、実装まとめにコードがあります。
動作保証はできませんが。
何番煎じだよというのはごめんなさい。
なるたけあいだの実装を省かないことで差別化しようと考えています。
1点更新、区間取得の最も基本的な構造にします。
遅延評価とかは気合があったら(か、必要に迫られたら)そのうち続編で書きます。
参考にしたサイトさん達
http://beet-aizu.hatenablog.com/entry/2017/09/10/132258
https://www.slideshare.net/iwiwi/ss-3578491
http://tsutaj.hatenablog.com/entry/2017/03/29/204841
セグメント木とは
何ができるの？
決まった性質を満たす操作、演算に関しての区間クエリを**O(logN)で求めることができる。(例えば1番目か５番目までの要素の最小値を求めるクエリ)
ただし、構築にはO(N)**かかるので注意
じゃあどんな演算なら使えるの？
モノイドと呼ばれるものなら埋め込める。
###モノイドとは
ちょっと厳密性を欠くと思うが、
結合則が成立し、単位元が存在すれば良い。
####結合則とは
例えば足し算で考えると、
$$(a+b)+c = a+(b+c)$$
となることで、3つ以上のものの間で演算するときにどこから計算しても結果が変わらなければ良い
これが成り立つから、セグメント木は値の取得で区間をどんどん分割していける。
([0,6)→([0,4)+[4,6)のように)
####単位元とは
元というと難しいが、応用上は、足し算で考えると、
$$a + 0 = 0+ a = a$$
となるように、それと演算すると、もとの数と同じになるような数を単位元と呼ぶ
他に、例えば掛け算を考えれば
$$ a \cdot 1 = 1 \cdot a = a$$
になって1が単位元である。
モノイド(にできるもの)の例
| 演算 | 単位元 |
|---|---|
| 足し算 | 0 |
| 掛け算 | 1 |
| 最小値 | INF ※1 |
| 最大値 | -INF ※2 |
| 最小公倍数 | 1 |
※1 要は出てきうる値以上の値をとれば単位元になる。例えば問題の制約で$A_i<=10^4$とかならば、単位元を$10^4+1$にすると、常に $min(A_i,10^4+1) = min(10^4+1,A_i) = A_i$となり単位元である。
全自然数とかが対象だと単位元は存在しない。
※2 最小値のときと似た議論で、出てきうる値以下のものをとると単位元になる。
他にも制約などによってはもっと特殊な演算もモノイドとして扱えることもあると思うが、そんなの思いつく人間はこんな記事読む必要もないと思うので置いておく。
[参考]
http://beet-aizu.hatenablog.com/entry/2017/09/10/132258
どんな構造なの？
カバーする区間に関してのクエリの結果を持つ2分木である。配列で実装する。
こんな感じ。下に行くほど区間がどんどん半々で小さくなっていく。
配列のインデックスに注目してほしいが、**親のインデックス×2,親のインデックス×2+1で子のインデックスが取得できる。**BIT(Binary index tree)とか二分木系では基本だけど初めて見たとき感動した。ちなみにこれは1-indexで考えているからで、0-indexのときは×2+1,×2+2になる。どっちでもいいといえばよいが、個人的に1-indexだと子から親を見る操作が、2でわって切り捨てるだけですむので1-indexにしている。
ここで一番下の葉にクエリをかけたい配列の元データが入ることになる。(今回だとサイズ4の配列である)
配列の要素数が2^kじゃなかったら？
一番下の葉の数が、もともと区間クエリを計算したい対象とする配列が入る分用意できるように木の深さを決める。
余った部分は単位元で埋めておけば良い。
さっきとすこし変えて配列サイズ=3の例を考えてみる。
余った部分に上で出てきた単位元を埋めておくと、
このように、特に影響を与えずに処理できる。
実装
初期化
末端の葉っぱの数が2^kの形で、クエリをかけたい配列のサイズ以上確保できるように、木を格納する配列サイズを決める。
$$1+2+4+8+\cdots+2^n = 2^{n+1} -1$$
であることを考えると、
$$2^k >= (データ配列サイズ)$$
となる、kについて、$2^{k+1}$のサイズの木格納用配列を用意すれば良い。
実装はこんな感じ。
class segtree:
def __init__(self, n,operator,identity):
"""
n:データ配列のサイズ
operator:演算子(モノイド)。関数オブジェクト
identity:演算子に対応する単位元
"""
nb = bin(n)[2:] #２進数に変換して戦闘の0bを取り除く
bc = sum([int(digit) for digit in nb]) #bitで1が立ってる数。これが1のときはちょうど2^nb.そうじゃないときは、2^nb<n<2^(nb+1)
if bc == 1: #2^nbなら
self.num_end_leaves = 2**(len(nb)-1) #最下段の葉っぱは2^nb個
else:#そうじゃないなら2^(nb+1)確保
self.num_end_leaves = 2**(len(nb))
self.array = [identity for i in range(self.num_end_leaves * 2)] #単位元で初期化
self.identity = identity
self.operator =operator
#後で使うので単位元と演算子を持っておく
値の更新
最初にちらっと書いたが、1-indexだと自分のインデックスを2で割って切り捨てると自分の親になる。(3→1,5→2のように)
これを大元までたどって順々に演算子を適用して行くことで値を更新する。
この部分だけの実装はこんな感じ。
def update(self,x,val):
"""
x:代入場所
val:代入する値
"""
actual_x = x+self.num_leaves #1-indexの末端の葉のindexがどこから始まるか分を足す(例えばデータ配列サイズ4のとき木配列サイズは8で、後半部はindex4から始まる。
self.array[actual_x] = val #値を更新する
while actual_x > 0 :
actual_x = actual_x//2#親を見る
self.array[actual_x] = self.operator(self.array[actual_x*2],self.array[actual_x*2+1])#あたらしい子をつかって親を更新
値の取得
ある範囲に関して、できるだけ大きい範囲をカバーする葉どうしの組み合わせでその範囲を表現できれば良い。
たとえば、データ配列のサイズが4で[0,2]の範囲が欲しかったとしたら、
こんな感じで、
- クエリに対して自分の担当範囲がはみ出していたら、子を見る。
- クエリの範囲に自分の担当範囲が入っていたら、値を返す。
- クエリと関係ない位置だったら単位元を返す。
を繰り返すと今回であれば[0,1],[2],単位元が値として得られてこれらの部分領域をマージしていく感じで、
$$operator([0,1],[2],単位元) = operator([0,1],[2]) = operator([0,2])$$
区間クエリを得ることができる。実装はこんなかんじ。
def get(self,q_left,q_right,arr_ind=1,leaf_left=0,depth=0):
"""
q_left：クエリ区間の左
q_right:クエリ区間の右
arr_ind:木配列のインデックス。最初は親なので1
leaf_left:木配列インデックスに対して、それが表す葉がカバーする範囲の左
depth:木配列での深さ。カバー範囲の広さの計算に使用
"""
width_of_floor = self.num_end_leaves//(2**depth) #今の葉のカバー幅
leaf_right = leaf_left+width_of_floor-1 #左端とカバー幅から今の葉のカバー範囲の右を求める。
if leaf_left > q_right or leaf_right < q_left:
return self.identity #クエリ領域と葉が関係ないなら単位元を返す
elif leaf_left >= q_left and leaf_right <= q_right:
return self.array[arr_ind] #クエリ領域に葉がすっぽり入ってるなら、葉の値を返す
else: #そうじゃないならば、子を見る
val_l = self.get(q_left,q_right,2*arr_ind,leaf_left,depth+1)#子の左
val_r = self.get(q_left,q_right,2*arr_ind+1,leaf_left+width_of_floor//2,depth+1)#子の右
return self.operator(val_l,val_r)#子をマージする演算をする。
#実装まとめ
今までのをくっつけただけなので、折りたたみ
class segtree:
def __init__(self, n,operator,identity):
nb = bin(n)[2:]
bc = sum([int(digit) for digit in nb])
if bc == 1:
self.num_end_leaves = 2**(len(nb)-1)
else:
self.num_end_leaves = 2**(len(nb))
self.array = [identity for i in range(self.num_end_leaves * 2)]
self.identity = identity
self.operator =operator
def update(self,x,val):
actual_x = x+self.num_end_leaves
self.array[actual_x] = val
while actual_x > 0 :
actual_x = actual_x//2
self.array[actual_x] = self.operator(self.array[actual_x*2],self.array[actual_x*2+1])
def get(self,q_left,q_right,arr_ind=1,leaf_left=0,depth=0):
width_of_floor = self.num_end_leaves//(2**depth)
leaf_right = leaf_left+width_of_floor-1
if leaf_left > q_right or leaf_right < q_left:
return self.identity
elif leaf_left >= q_left and leaf_right <= q_right:
return self.array[arr_ind]
else:
val_l = self.get(q_left,q_right,2*arr_ind,leaf_left,depth+1)
val_r = self.get(q_left,q_right,2*arr_ind+1,leaf_left+width_of_floor//2,depth+1)
return self.operator(val_l,val_r)
お試し
適当なrangeの配列を作って、最小値を聞いてみる。
ついでに実行時間も適当に図る。
s_tree = segtree(10**5,min,10**9) #10**5までの配列なので適当に単位元は10**9
arr = [i for i in range(10**5)]
print(datetime.datetime.now())
for i,a in enumerate(arr):
s_tree.update(i,a)
print(datetime.datetime.now())
print(s_tree.get(0,10**4))
print(s_tree.get(3*10**4,5*10**4))
print(s_tree.get(2,7*10**4))
print(datetime.datetime.now())
2020-02-04 19:15:34.934814
2020-02-04 19:15:35.646339
0
30000
2
2020-02-04 19:15:35.646587
やっぱり構築が一番重い。
動作は正常そう。(もっと複雑な問題で動かしてみないとバグ埋まってそうですが。)