leetcode-go/hard/699.falling-squares.go

type node struct {
	l     int
	r     int
	left  *node
	right *node
	max   int
	tag   int // Lazy
}

var tree []node
var tid int

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]
	root.left.buildTree(l, (l+r)/2)
	root.right.buildTree((l+r)/2+1, r)
}

func (root *node) update(i, j, val int) {
	if root.l == i && root.r == j {
		root.tag = val
		return
	}
	root.max = maxInt(root.max, val)
	if m := (root.l + root.r) / 2; j <= m {
		root.left.update(i, j, val) // Update left tree
	} else if m+1 <= i {
		root.right.update(i, j, val) // Right tree
	} else {
		root.left.update(i, m, val) // Split into 2 parts
		root.right.update(m+1, j, val)
	}
}

func (root *node) query(i, j int) int {
	if root.l == i && root.r == j {
		if root.tag != 0 {
			return root.tag
		}
		return root.max
	}
	m := (root.l + root.r) / 2
	if root.tag != 0 {
		root.max = root.tag // PushDown method
		root.left.update(root.l, m, root.tag)
		root.right.update(m+1, root.r, root.tag)
	}
	root.tag = 0 // Finished update
	if j <= m {
		return root.left.query(i, j)
	} else if m+1 <= i {
		return root.right.query(i, j)
	} else {
		return maxInt(root.left.query(i, m), root.right.query(m+1, j))
	}
}

func fallingSquares(positions [][]int) []int {
	pts := make([]int, 0)
	m := make(map[int]int)
	for _, pos := range positions {
		beg, end := pos[0], pos[0]+pos[1]-1
		if _, ok := m[beg]; !ok {
			m[beg] = 0
			pts = append(pts, beg)
		}
		if _, ok := m[end]; !ok {
			m[end] = 0
			pts = append(pts, end)
		}
	}
	sort.Ints(pts)
	for i := range pts {
		m[pts[i]] = i + 1
	} // Discretization
	n := len(pts)
	tree = make([]node, 2*n)
	tid = 0
	root := &tree[0]
	root.buildTree(1, n)
	res := make([]int, len(positions))
	for idx, pos := range positions {
		i, j := m[pos[0]], m[pos[0]+pos[1]-1]
		max := root.query(i, j)
		root.update(i, j, max+pos[1])
		res[idx] = root.query(1, n)
	}
	return res
}

func maxInt(x, y int) int {
	if x < y {
		return y
	}
	return x
}