はじめに
普段はググってセグ木を貼るだけでしたが、コロナウイルスの影響で時間が出来たので原理を理解しながら実装してみました。
ネットを中心に様々な方の記事を参考にしながら構築していきました。最後の参考文献にリンク先は載せますが、普段コンテスト中にコードを利用させていただいているじゅっぴーさんの記事、ライブラリをよくのぞかせていただいてるいかたこのたこつぼさんのHPに特にお世話になったので、冒頭で載せておきます。
目的
原理などは別の記事にまとめようと思います。この記事の目的は自分のセグ木の使い方の説明です。
おおまかな特徴としては、じゅっぴーさんを参考にさせていただいて、自分で単位元と行いたい操作を設定すること、クラスで書かれていることが挙げられます。セグ木のノードは1-indexで書きました。
コード
#####segfunc#####
def segfunc(x, y):
return
#################
#####ide_ele#####
ide_ele =
#################
class SegTree:
"""
init(init_val, ide_ele): 配列init_valで初期化 O(N)
update(k, x): k番目の値をxに更新 O(logN)
query(l, r): 区間[l, r)をsegfuncしたものを返す O(logN)
"""
def __init__(self, init_val, segfunc, ide_ele):
"""
init_val: 配列の初期値
segfunc: 区間にしたい操作
ide_ele: 単位元
n: 要素数
num: n以上の最小の2のべき乗
tree: セグメント木(1-index)
"""
n = len(init_val)
self.segfunc = segfunc
self.ide_ele = ide_ele
self.num = 1 << (n - 1).bit_length()
self.tree = [ide_ele] * 2 * self.num
# 配列の値を葉にセット
for i in range(n):
self.tree[self.num + i] = init_val[i]
# 構築していく
for i in range(self.num - 1, 0, -1):
self.tree[i] = self.segfunc(self.tree[2 * i], self.tree[2 * i + 1])
def update(self, k, x):
"""
k番目の値をxに更新
k: index(0-index)
x: update value
"""
k += self.num
self.tree[k] = x
while k > 1:
self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
k >>= 1
def query(self, l, r):
"""
[l, r)のsegfuncしたものを得る
l: index(0-index)
r: index(0-index)
"""
res = self.ide_ele
l += self.num
r += self.num
while l < r:
if l & 1:
res = self.segfunc(res, self.tree[l])
l += 1
if r & 1:
res = self.segfunc(res, self.tree[r - 1])
l >>= 1
r >>= 1
return res
使い方
1. 初期化用のリストを用意する
なんでもいいです。例えば↓
a = [14, 5, 9, 13, 7, 12, 11, 1, 7, 8]
2. 区間に行う操作を決める
区間の最小値、最大値を求めたい、区間の和を求めたいなど。今回は、最小値を求めるとします。segfuncに関数を書き込んでください。
def segfunc(x, y):
return min(x, y)
3. 単位元を定める
初期化に使います。演算に影響を与えないものです。最小値を求めるなら無限大がこれに当たります。
ide_ele = float('inf')
4. オブジェクトを作成、引数は上の3つ
seg = SegTree(a, segfunc, ide_ele)
5. 各操作を行えます
リストは0-indexで扱ってください。（いつも通りです。）
行える操作は以下の2つです。
1. ある1つの要素の更新
update(k, x) : k番目の要素をxに更新します。
2. ある区間の操作の結果を取得
query(l, r) : [l, r)(l以上r未満の区間)から値を取得します。
# [0, 8)の最小値を表示
print(seg.query(0, 8)) # 1
# 5番目を0に変更
seg.update(5, 0)
# [0, 8)の最小値を表示
print(seg.query(0, 8)) # 0
まとめ
長いので折り込みました。
#####segfunc#####
def segfunc(x, y):
return min(x, y)
#################
#####ide_ele#####
ide_ele = float('inf')
#################
class SegTree:
"""
init(init_val, ide_ele): 配列init_valで初期化 O(N)
update(k, x): k番目の値をxに更新 O(N)
query(l, r): 区間[l, r)をsegfuncしたものを返す O(logN)
"""
def __init__(self, init_val, segfunc, ide_ele):
"""
init_val: 配列の初期値
segfunc: 区間にしたい操作
ide_ele: 単位元
n: 要素数
num: n以上の最小の2のべき乗
tree: セグメント木(1-index)
"""
n = len(init_val)
self.segfunc = segfunc
self.ide_ele = ide_ele
self.num = 1 << (n - 1).bit_length()
self.tree = [ide_ele] * 2 * self.num
# 配列の値を葉にセット
for i in range(n):
self.tree[self.num + i] = init_val[i]
# 構築していく
for i in range(self.num - 1, 0, -1):
self.tree[i] = self.segfunc(self.tree[2 * i], self.tree[2 * i + 1])
def update(self, k, x):
"""
k番目の値をxに更新
k: index(0-index)
x: update value
"""
k += self.num
self.tree[k] = x
while k > 1:
self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
k >>= 1
def query(self, l, r):
"""
[l, r)のsegfuncしたものを得る
l: index(0-index)
r: index(0-index)
"""
res = self.ide_ele
l += self.num
r += self.num
while l < r:
if l & 1:
res = self.segfunc(res, self.tree[l])
l += 1
if r & 1:
res = self.segfunc(res, self.tree[r - 1])
l >>= 1
r >>= 1
return res
a = [14, 5, 9, 13, 7, 12, 11, 1, 7, 8]
seg = SegTree(a, segfunc, ide_ele)
print(seg.query(0, 8))
seg.update(5, 0)
print(seg.query(0, 8))
その他の操作（最小値以外）
| 操作 | segfunc | 単位元 |
|---|---|---|
| 最小値 | min(x, y) | float('inf') |
| 最大値 | max(x, y) | -float('inf') |
| 区間和 | x + y | 0 |
| 区間積 | x * y | 1 |
| 最大公約数 | math.gcd(x, y) | 0 |
おわりに
間違いやアドバイスなどがあったら教えてください。使用する際は十分テストしてからコンテストに使ってください!!
参考文献
冒頭で述べたお二方です。↓
・じゅっぴーさんのブログ
・いかたこのたこつぼさんのHP
原理の全般の理解に助かりました。↓
・セグメント木を一歩一歩実装して理解しようとする(python)
非再帰セグ木の部分の参考になりました。↓
・Segment Treeをちょっと高速化したい
・Segment Treeを非再帰で書く