Goでセグメント木を作る
初めに
Goで競技プログラミングを戦うみなさんは先日のABC185F問題は解けましたか？
どうも私の知人がセグメント木だとは気付けたものの諦めてしまったそうです。いわゆる 貼るだけ 問題なので残念ですね。
AtCoder Library (ACL) が公開されている以上、C++以外を選択するならば自作する覚悟が必要でしょうか。この記事では蟻本を参考にしつつGoでの実装をしようと思います。
※筆者はC++で参加している。
木を表現する
セグメント木とは、区間に対する問い合わせと1要素の更新をO(logN)で行える完全二分木（全ての葉の深さが等しい木）です。
木構造というと下記のような構造を想像するでしょうか？
type struct Node {
value int
right *Node
left *Node
}
しかしここでは配列を用いて木を表現してみましょう。ヒープを自作したことがあれば馴染みがあるかもしれません。
- 根の添字を1とする。
- 添字がnである頂点の左の子の添字は2n
- 添字がnである頂点の右の子の添字は2n+1
言うまでもなく添字がnの頂点の親の添字はn/2です。根の添字を0として、子をそれぞれ2n+1と2n+2にしても構いません。その場合の親は(n-1)/2です。蟻本が0originなので、ここでは根を0とします。
完全二分木なので葉の個数は2のべき乗で、頂点数は2^0+2^1+2^2+...+2^dです。配列の長さは扱いたい区間の長さを超える最小の2のべき乗としておきましょう。
type SegmentTree struct {
data []int
n int
}
func NewSegmentTree(n int) *SegmentTree {
segtree := new(SegmentTree)
segtree.n = 1
for segtree.n < n {
segtree.n *= 2
}
segtree.data = make([]int, segtree.n*2-1)
return segtree
}
このsegtree.data[i]が頂点iの表す区間に対する計算結果だとします。
頂点と区間
根は扱いたい区間の全体を表します。2つの子はそれぞれその半分の区間を表すものとします。
つまり、ある頂点の表す区間が [a, b) であるならその子の表す区間はそれぞれ、 [a, (a+b)/2) と [(a+b)/2, b) です。
※ここでは区間は全て0originの半開く区間とします。
func (segtree *SegmentTree) query(begin, end, idx, a, b int) {
if b <= begin || end <= a {
// クエリと関係のない区間
return
}
if begin <= a && b <= end {
// 全体がクエリの対象になる区間
return
}
// 一部がクエリの対象にならない場合は子に尋ねる
v1 := segtree.query(begin, end, idx*2+1, a, (a+b)/2)
v2 := segtree.query(begin, end, idx*2+2, (a+b)/2, b)
// DO SOMETHING
}
func (segtree *SegmentTree) Query(begin, end int) {
segtree.query(begin, end, 0, 0, segtree.n)
}
どういうときに必要？
木をどう表現するか分かったところでどのような問題にセグメント木が必要か考えてみましょう。最も簡単な例は累積和でしょうか。
長さNの数列が与えられたとき、ある区間に含まれる数の総和を求めるよくあるアレです。iからjまでの総和を求めるためには添字0からの和を事前に計算しておけば0からjまでの和から0からiまでの和を引き算すれば、構築O(N)クエリO(1)になります。。
type PrefixSum struct {
sum []int64
}
func NewPrefixSum(values []int64) *PrefixSum {
ps := &PrefixSum{}
ps.sum = append(acc.sum, 0)
var sum int64
for _, v := range values {
sum += v
ps.sum = append(acc.sum, sum)
}
return ps
}
func (ps *PrefixSum) Query(begin, end int) int64 {
return ps.sum[end] - ps.sum[begin]
}
しかし、この方法では更新（例えばx番目の要素をyに置き換えたい）が構築と同じO(N)になってしまいます。1要素の更新であればO(logN)で行えます。区間を扱いそうな問題に直面した場合、本当に区間に対する問い合わせと1要素の更新で済むか確認しましょう。
クエリ
累積和を例に挙げたのでひとまず区間に対する問い合わせは総和だとして考えていきます。
この場合はsegtree.data[i]が頂点iの区間の総和になります。もちろんsegtree.data[i] = segtree.data[i*2+1] + segtree.data[i*2+2] です。
そして、ある頂点の表す区間がクエリにどのように含まれるかによって木の辿り方が変わります。
- ある頂点の表す区間がクエリに完全に含まれるならば、その区間の総和を足す。
- ある頂点の表す区間がクエリと交差しないならば総和を0とする。
- ある頂点の表す区間がクエリと一部だけ交差するなら、完全に含まれる区間を探して子に尋ねる。
func (segtree *SegmentTree) query(begin, end, idx, a, b int) int {
if b <= begin || end <= a {
// クエリと関係のない区間
return 0
}
if begin <= a && b <= end {
// 全体がクエリの対象になる区間
return segtree.data[idx]
}
// 一部がクエリの対象にならない場合は子に尋ねる
v1 := segtree.query(begin, end, idx*2+1, a, (a+b)/2)
v2 := segtree.query(begin, end, idx*2+2, (a+b)/2, b)
return v1 + v2
}
func (segtree *SegmentTree) Query(begin, end int) int {
return segtree.query(begin, end, 0, 0, segtree.n)
}
更新
親は (n-1)/2 もしくは n/2 でした。そして segtree.data[i] = segtree.data[i*2+1] + segtree.data[i*2+2] です。
長さ1の区間（つまり葉）の値を書き換えて、それを根に向かって更新していくことを考えます。
func (segtree *SegmentTree) Update(idx, x int) {
idx += segtree.n - 1
segtree.data[idx] = x
for 0 < idx {
idx = (idx - 1) / 2
segtree.data[idx] = segtree.data[idx*2+1] + segtree.data[idx*2+2]
}
}
ここで説明が必要になる点は idx += segtree.n - 1 でしょう。
+---------------+
| 0 |
+-------+-------+
| 1 | 2 |
+---+---+-------+
| 3 | 4 | 5 | 6 |
+-+-+---+-+-+-+-+
|7|8|9|A|B|C|D|E|
+-+-+-+-+-+-+-+-+
アスキーアートで申し訳ないのですが、下段が葉、上段が根の区間だとします。16進数で添字を書いてみました。我ながら分かりづらい。
深さdの頂点は2^d個存在します。もし葉が深さnならば葉より浅い頂点の個数は2^n-1個存在することになります。
もとの配列でi番目の要素を書き換えたい場合、i+2^n-1 が葉の要素の添字になっていています。これを根に向かって更新しているわけです。
ABC084D
とりあえず累積和の問題を探したらABC084Dが見つかりました。prefix sumでなくセグメント木で解いてみましょう。なお、入出力についてはこちらの記事が便利です。
iが 2017に似た数 ならばi番目の要素を1に、そうでないなら0を要素に持つような長さ10^5の配列の区間の和を求める問題です。
和以外
累積和を扱いましたがABC185Fはどのように解けば良いのでしょう。この問題は排他的論理和です。
少し回り道をして、左の子は自身より小さい値を持つ、右の子は自身より大きな値を持つ、というようなよくある2分木を思い浮かべて下さい。こういった木構造は便利ですし皆さん恩恵を受けていると思います。では数値以外の場合はどのような値がこのような木構造で扱うことが出来るのでしょうか？実は制約があり、大小比較が可能である必要があります。
次にバケットソートを思い浮かべてみます。
func main() {
const N = 1000
cnt := make([]int, N)
var n int
fmt.Scan(&n)
for i := 0; i < n; i++ {
var m int
fmt.Scan(&m)
cnt[m]++
}
for i := 0; i < N; i++ {
for j := 0; j < cnt[i]; j++ {
fmt.Println(i)
}
}
}
バケットソートは入力が整数値で値の上限があまり大きくない場合に利用できるO(n)のソートアルゴリズムです。
2分木やバケットソートから何を言いたいかというと、扱う値の特徴によってアルゴリズムやデータ構造は最適化出来るし、特徴次第では使えないものが存在するということです。
モノイド
セグメント木で扱えるデータはモノイドです。
集合（型）と2項演算子の組で、結合律があって単位元が存在するようなものです。
累積和の場合を考えると、a+(b+c) = (a+b)+c ですし単位元は0（x+0=x）です。
minやmaxも同様で、 min(a, min(b, c)) = min(min(a, b), c) ですし min(x,∞)=x です。
ではABC185Fの場合はどうでしょう。これは整数値と排他的論理和でした。a⊕(b⊕c) = (a⊕b)⊕c で x⊕0=x ですね？
そもそも問題文がセグメント木を感じさせます。
区間への問い合わせ、1要素の更新という2点に注意するように書きましたが、実はそれだけでなく、扱いたいものがモノイドであるかにも注目して下さい。
2項演算子や単位元も構造体に追加して以下のようにしましょう。ジェネリクスのあるバージョンのGoがAtCoderにやってくるまでは何もかもinterface{}です。2のべき乗にした分は単位元で初期化します。
type SegmentTree struct {
data []interface{}
n int
e interface{}
op func(interface{}, interface{}) interface{}
}
func NewSegmentTree(n int, e interface{}, op func(interface{}, interface{}) interface{}) *SegmentTree {
segtree := new(SegmentTree)
segtree.e = e
segtree.op = op
segtree.n = 1
for segtree.n < n {
segtree.n *= 2
}
segtree.data = make([]interface{}, segtree.n*2-1)
for i := 0; i < segtree.n*2-1; i++ {
segtree.data[i] = segtree.e
}
return segtree
}
func (segtree *SegmentTree) Update(idx int, x interface{}) {
idx += segtree.n - 1
segtree.data[idx] = x
for 0 < idx {
idx = (idx - 1) / 2
segtree.data[idx] = segtree.op(segtree.data[idx*2+1], segtree.data[idx*2+2])
}
}
func (segtree *SegmentTree) query(begin, end, idx, a, b int) interface{} {
if b <= begin || end <= a {
return segtree.e
}
if begin <= a && b <= end {
return segtree.data[idx]
}
v1 := segtree.query(begin, end, idx*2+1, a, (a+b)/2)
v2 := segtree.query(begin, end, idx*2+2, (a+b)/2, b)
return segtree.op(v1, v2)
}
func (segtree *SegmentTree) Query(begin, end int) interface{} {
return segtree.query(begin, end, 0, 0, segtree.n)
}
// NewSegmentTree(N, 0, func(a, b interface{}) interface{} { return a.(int) ^ b.(int) })
ACしたGoの解答 ※ fmt.ScanfするとTLEする。
その先へ
ACLには、BIT(Fenwick Tree)やLazy Segtreeも実装されています。区間をより上手に扱いたいGo競技プログラマーの皆さんは次にそれらへの理解を深めてみましょう。
おわりに
何かのアドベントカレンダーの記事にすれば良かった。
Discussion