package main

import "fmt"

type SegmentTree []int

// 获取左孩子下标
func (st SegmentTree) Left(i int) int {
	v := i * 2 + 1
	if v >= len(st) {
		return -1
	}
	return v
}

// 获取右孩子下标
func (st SegmentTree) Right(i int) int {
	v := i * 2 + 2
	if v >= len(st) {
		return -1
	}
	return v
}

// 更新线段树
func (st SegmentTree) change(node int, index int, value int, s, e int) {
	if s == e {
		st[node] = value
		return
	}
	mid := s + (e-s)/2
	// 按照范围区间选择左右子树
	if index <= mid {
		st.change(st.Left(node), index, value, s, mid)
	} else {
		st.change(st.Right(node), index, value, mid + 1, e)
	}
	st[node] = st[st.Left(node)] + st[st.Right(node)]
}

// 范围查询
func (st SegmentTree) range_(node int, left, right int, s, e int) int {
	// 如果查询范围与当前的遍历范围匹配，则直接返回
	if left == s && right == e {
		return st[node]
	}

	mid := s + (e-s)/2
	// 若查询范围位于mid左侧（包含mid），则只需要遍历左子树
	if right <= mid {
		return st.range_(st.Left(node), left, right, s, mid)
	}
	// 若查询范围位于mid右侧（不包含mid），则只需要遍历右子树
	if left > mid {
		return st.range_(st.Right(node), left, right, mid + 1, e)
	}
	// 分段遍历，将结果相加后返回
	return st.range_(st.Left(node), left, mid, s, mid) + st.range_(st.Right(node), mid + 1, right, mid + 1, e)
}

func (st SegmentTree) build(nums []int, node int, s, e int) {
	if s == e {
		st[node] = nums[s]
		return
	}
	mid := s + (e-s)/2
	left := st.Left(node)
	right := st.Right(node)
	st.build(nums, left, s, mid)
	st.build(nums, right, mid + 1, e)
	st[node] = st[left] + st[right]
}

func (st SegmentTree) Update(index int, value int) {
	st.change(0, index, value, 0, len(st)/4-1)
}

func (st SegmentTree) Range(left, right int) int {
	return st.range_(0, left, right, 0, len(st)/4-1)
}

func (st SegmentTree) LevelTravel() {
	fmt.Println("层序遍历线段树: ")
	q := make([]int, 0)
	q = append(q, 0)
	depth := 0
	for len(q) != 0 {
		length := len(q)
		fmt.Printf("%d: ", depth)
		for ; length > 0; length-- {
			head := q[0]
			fmt.Printf("%d ", st[head])
			q = q[1:]
			if st.Left(head) != -1 {
				q = append(q, st.Left(head))
			}
			if st.Right(head) != -1 {
				q = append(q, st.Right(head))
			}
		}
		depth++
		fmt.Println()
	}
	fmt.Println()
}

func InitializeSegmentTree(nums []int) SegmentTree {
	n := len(nums)
	st := SegmentTree{}
	st = make([]int, 4 * n)
	st.build(nums, 0, 0, n - 1)
	return st
}

func main() {
	nums := []int{1,2,3,4,5,6,7}
	fmt.Println("初始化数组:", nums, "\n")
	st := InitializeSegmentTree(nums)
	st.LevelTravel()
	fmt.Println("[0,3]范围和:", st.Range(0, 3))
	fmt.Println("修改[1]为20后: ")
	st.Update(1, 20)
	st.LevelTravel()
	fmt.Println("[0,3]范围和:", st.Range(0, 3))
}