leetcode-go/basis/algorithms/segmenttree/segmenttree.go

package main

import (
	"fmt"
)

type node struct {
	l     int
	r     int
	left  *node
	right *node
	val   int
	inc   int // Lazy
}

func (n node) mid() int {
	return n.l + (n.r-n.l)/2
}

var tree []node
var tid int

func newTree(n int) *node {
	tree = make([]node, 2*n)
	tid = 0
	root := &tree[0]
	root.buildTree(1, n)
	return root
}

func (root *node) buildTree(l, r int) {
	root.l = l
	root.r = r
	if l == r {
		return
	}
	tid++
	root.left = &tree[tid]
	tid++
	root.right = &tree[tid]
	m := l + (r-l)/2
	root.left.buildTree(l, m)
	root.right.buildTree(m+1, r)
}

func (root *node) update(l, r, val int) {
	if root.l == l && root.r == r {
		root.inc += val
		return
	}
	root.val += (r - l + 1) * val
	if r <= root.mid() {
		root.left.update(l, r, val)
	} else if root.mid()+1 <= l {
		root.right.update(l, r, val)
	} else {
		root.left.update(l, root.mid(), val)
		root.right.update(root.mid()+1, r, val)
	}
}

func (root *node) query(l, r int) int {
	if root.l == l && root.r == r {
		return root.val + (r-l+1)*root.inc
	}
	if root.inc != 0 {
		root.val += (root.r - root.l + 1) * root.inc   // PushDown
		root.left.update(root.l, root.mid(), root.inc) // !!the RANGE is important
		root.right.update(root.mid()+1, root.r, root.inc)
	}
	root.inc = 0 // Finished update
	if r <= root.mid() {
		return root.left.query(l, r)
	} else if root.mid()+1 <= l {
		return root.right.query(l, r)
	} else {
		return root.left.query(l, root.mid()) + root.right.query(root.mid()+1, r)
	}
}

func main() {
	var n int
	fmt.Scan(&n)
	root := newTree(n)
	for i := 1; i <= n; i++ {
		root.update(i, i, i)
	}
	fmt.Println("Sum of 1 ~", n, "is", root.query(1, n))
	var l, r, val int
	fmt.Print("Update(l, r, val): ")
	fmt.Scan(&l, &r, &val)
	root.update(l, r, val)
	fmt.Print("Query(l, r): ")
	fmt.Scan(&l, &r)
	fmt.Println(root.query(l, r))
}