Goでセグ木を実装する
はじめに
競プロでセグ木を使わないと厳しい問題がちらほらと出てきたので、セグ木を実装していきたいと思います。
c++で書かれているサンプルや解説はあるので、そちらを参照されるといいと思います。ちなみに、このサイトのものをGo用に書き換えただけです。
今回実装するセグ木は、普通のセグ木と遅延評価セグ木です。
他のやつはそのうちやりたいと思います。
何ができるか
まず、セグ木について、何ができるかというと、元となる要素n個の数列があったとすると、
1. 1要素の変更をO(logN)でできる。
2. [a,b)区間上の最小値や合計値をO(logNで)取得できる。(今回は、最小値)
次に、遅延評価付きセグ木は何ができるかというと、
1. [a,b)区間の要素に対する変更をO(logN)でできる。(今回は一律加算)
2. [a,b)区間の最小値や合計をO(logN)で取得できる。(今回は合計)
実装(セグ木)
標準入力周りで、若干汚くなってますが、最初に数列を用意してそれに対し操作をしていきます。
seg.go
package main
import (
"bufio"
"fmt"
"os"
"strconv"
)
var sc = bufio.NewScanner(os.Stdin)
func iScan() int {
n, _ := strconv.Atoi(Scan())
return n
}
func larger(a, b int) int {
if a < b {
return b
} else {
return a
}
}
func smaller(a, b int) int {
if a > b {
return b
} else {
return a
}
}
var mod int = 1000000007
var inf int = 1000000000000
type segtree struct {
n int
node []int
}
func newSegtree(v []int) *segtree {
sz := len(v)
_n := 1
for _n < sz {
_n *= 2
}
_node := make([]int, 2*_n-1)
for i := 0; i < _n; i++ {
_node[i] = inf
}
for i := 0; i < sz; i++ {
_node[i+_n-1] = v[i]
}
for i := _n - 2; i >= 0; i-- {
_node[i] = smaller(_node[2*i+1], _node[2*i+2])
}
return &segtree{
n: _n,
node: _node,
}
}
func (u segtree) update(x int, val int) {
x += u.n - 1
u.node[x] = val
for x > 0 {
x = (x - 1) / 2
u.node[x] = smaller(u.node[2*x+1], u.node[2*x+2])
}
return
}
func (u segtree) getmin(a int, b int, k int, l int, r int) int {
if r < 0 {
r = u.n
}
if r <= a || b <= l {
return inf
}
if a <= l && r <= b {
return u.node[k]
}
vl := u.getmin(a, b, 2*k+1, l, (l+r)/2)
vr := u.getmin(a, b, 2*k+2, (l+r)/2, r)
return smaller(vl, vr)
}
func main() {
//標準入力用
buf := make([]byte, 0)
sc.Buffer(buf, mod)
sc.Split(bufio.ScanWords)
//標準入力
n := iScan()
a := make([]int, n)
for i := 0; i < n; i++ {
a[i] = iScan()
}
u := newSegtree(a)
fmt.Println(u)
for i := 0; i < 3; i++ {
x := iScan()
val := iScan()
// x番目の値を valに変更
u.update(x, val)
b := iScan()
c := iScan()
// [b,c)内の最小値の取得
fmt.Println(u.getmin(b, c, 0, 0, -1))
fmt.Println(u.node)
}
}
実装(遅延評価セグ木)
lazy.go
package main
import (
"bufio"
"fmt"
"os"
"strconv"
)
var sc = bufio.NewScanner(os.Stdin)
func iScan() int {
n, _ := strconv.Atoi(Scan())
return n
}
func larger(a, b int) int {
if a < b {
return b
} else {
return a
}
}
func smaller(a, b int) int {
if a > b {
return b
} else {
return a
}
}
var mod int = 1000000007
var inf int = 1000000000000
type lazysegtree struct {
n int
node []int
lazy []int
}
func newlazySegtree(v []int) *lazysegtree {
sz := len(v)
_n := 1
for _n < sz {
_n *= 2
}
_node := make([]int, 2*_n-1)
_lazy := make([]int, 2*_n-1)
for i := 0; i < sz; i++ {
_node[i+_n-1] = v[i]
}
for i := _n - 2; i >= 0; i-- {
_node[i] = _node[2*i+1] + _node[2*i+2]
}
return &lazysegtree{
n: _n,
node: _node,
lazy: _lazy,
}
}
// k番目のノードについて遅延評価を行う
func (u lazysegtree) eval(k int, l int, r int) {
if u.lazy[k] != 0 {
u.node[k] += u.lazy[k]
if r-l > 1 {
u.lazy[2*k+1] += u.lazy[k] / 2
u.lazy[2*k+2] += u.lazy[k] / 2
}
u.lazy[k] = 0
}
return
}
//[a,b)に対して加算
func (u lazysegtree) add(a int, b int, x int, k int, l int, r int) {
if r < 0 {
r = u.n
}
u.eval(k, l, r)
if b <= l || r <= a {
return
}
if a <= l && r <= b {
u.lazy[k] += (r - l) * x
u.eval(k, l, r)
} else {
u.add(a, b, x, 2*k+1, l, (l+r)/2)
u.add(a, b, x, 2*k+2, (l+r)/2, r)
u.node[k] = u.node[2*k+1] + u.node[2*k+2]
}
}
//[a,b)内の合計の取得
func (u lazysegtree) getsum(a int, b int, k int, l int, r int) int {
if r < 0 {
r = u.n
}
if r <= a || b <= l {
return 0
}
u.eval(k, l, r)
if a <= l && r <= b {
return u.node[k]
}
vl := u.getsum(a, b, 2*k+1, l, (l+r)/2)
vr := u.getsum(a, b, 2*k+2, (l+r)/2, r)
return vl + vr
}
func main() {
buf := make([]byte, 0)
sc.Buffer(buf, mod)
sc.Split(bufio.ScanWords)
//予備
n := iScan()
a := make([]int, n)
for i := 0; i < n; i++ {
a[i] = iScan()
}
u := newlazySegtree(a)
fmt.Println(u)
for i := 0; i < 2; i++ {
a := iScan()
b := iScan()
x := iScan()
u.add(a, b, x, 0, 0, -1)
fmt.Println(u)
}
}
最後に
他にも種類があるようなので、実装し次第あげていきたいと思います。