segment_tree.go

package main

import "fmt"

type SegmentTree struct {
	offset      int
	inf         interface{}
	data        []interface{}
	merge       Merger
}

type Merger func(a, b interface{}) interface{}

func IntMergerMin(a, b interface{}) interface{} {
	aInt := a.(int)
	bInt := b.(int)
	if aInt < bInt {
		return aInt
	}
	return bInt
}

func InitIntSegmentTree(a []int, inf int, merge Merger) *SegmentTree {
	n := len(a)
	vec := make([]interface{}, n)
	for i := 0 ; i < n ; i++ {
		vec[i] = interface{}(a[i])
	}
	return InitSegmentTree(vec, inf, merge)
}

func InitSegmentTree(a []interface{}, inf interface{}, merge Merger) *SegmentTree {
	n := len(a)
	size := 1
	for size < n {
		size *= 2
	}
	data := make([]interface{}, size*2)
	for j, i := 0, size-1 ; i < size+n-1 ; i++ {
		data[i] = a[j]
		j++
	}
	for i := size+n-1 ; i < size*2 ; i++ {
		data[i] = inf
	}

	for i := size-2 ; i >= 0 ; i-- {
		data[i] = merge(data[i*2+1], data[i*2+2])
	}

	return &SegmentTree{
		inf:    inf,
		offset: size,
		data:   data,
		merge:  merge,
	}
}

func (tree *SegmentTree) RangeMin(from, to int) interface{} {
	return tree.rangeMin(from, to, 0, 0, tree.offset)
}

func (tree *SegmentTree) rangeMin(from, to, index, left, right int) interface{} {
	if to <= left || right <= from {
		return tree.inf
	}
	if from <= left && right <= to {
		return tree.data[index]
	}
	med := (left+right)/2
	lvalue := tree.rangeMin(from, to, index*2+1, left, med)
	rvalue := tree.rangeMin(from, to, index*2+2, med, right)

	return tree.merge(lvalue, rvalue)
}

func (tree *SegmentTree) UpdateAt(index int, value interface{}) {
	idx := tree.offset-1+index
	tree.data[idx] = value
	for idx >= 1 {
		parent := idx / 2
		left := parent*2
		right := parent*2+1
		tree.data[parent] = tree.merge(tree.data[left], tree.data[right])
		idx = parent
	}
}

func main() {
	seg := InitIntSegmentTree([]int{1, 9, 5, 3, 7, 2, 4, 6, 8}, 1e9, IntMergerMin)

	fmt.Println(seg.RangeMin(1, 2))
	fmt.Println(seg.RangeMin(3, 5))
	fmt.Println(seg.RangeMin(0, 10))
	fmt.Println(seg.RangeMin(5, 7))
	fmt.Println(seg.RangeMin(2, 4))


	seg.UpdateAt(2, 1)
	fmt.Println(seg.RangeMin(1, 3))
}