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 }