Codeforces 916E — Jamie and Tree

1package main2 3import (4	"bufio"5	. "fmt"6	"io"7	"math/bits"8)9 10// https://space.bilibili.com/20621411type fenwickDiff16 [][2]int12 13func (t fenwickDiff16) _add(i, val int) {14	for iv := i * val; i < len(t); i += i & -i {15		t[i][0] += val16		t[i][1] += iv17	}18}19 20func (t fenwickDiff16) add(l, r, val int) {21	t._add(l, val)22	t._add(r+1, -val)23}24 25func (t fenwickDiff16) pre(i0 int) int {26	var s0, s1 int27	for i := i0; i > 0; i &= i - 1 {28		s0 += t[i][0]29		s1 += t[i][1]30	}31	return (i0+1)*s0 - s132}33 34func (t fenwickDiff16) query(l, r int) int {35	return t.pre(r) - t.pre(l-1)36}37 38func cf916E(_r io.Reader, _w io.Writer) {39	in := bufio.NewReader(_r)40	out := bufio.NewWriter(_w)41	defer out.Flush()42 43	var n, q, dfn, op, v, w, val, rt int44	Fscan(in, &n, &q)45	a := make([]int, n)46	for i := range a {47		Fscan(in, &a[i])48	}49	g := make([][]int, n)50	for i := 1; i < n; i++ {51		Fscan(in, &v, &w)52		v--53		w--54		g[v] = append(g[v], w)55		g[w] = append(g[w], v)56	}57 58	nodes := make([]struct{ l, r int }, n)59	const mx = 1760	pa := make([][mx]int, n)61	dep := make([]int, n)62	var build func(int, int) int63	build = func(v, p int) int {64		dfn++65		nodes[v].l = dfn66		pa[v][0] = p67		sz := 168		for _, w := range g[v] {69			if w != p {70				dep[w] = dep[v] + 171				sz += build(w, v)72				a[v] += a[w]73			}74		}75		nodes[v].r = nodes[v].l + sz - 176		return sz77	}78	build(0, -1)79	for i := 0; i+1 < mx; i++ {80		for v := range pa {81			if p := pa[v][i]; p != -1 {82				pa[v][i+1] = pa[p][i]83			} else {84				pa[v][i+1] = -185			}86		}87	}88	uptoDep := func(v, d int) int {89		for k := uint(dep[v] - d); k > 0; k &= k - 1 {90			v = pa[v][bits.TrailingZeros(k)]91		}92		return v93	}94	getLCA := func(v, w int) int {95		if dep[v] > dep[w] {96			v, w = w, v97		}98		w = uptoDep(w, dep[v])99		if w == v {100			return v101		}102		for i := mx - 1; i >= 0; i-- {103			if pv, pw := pa[v][i], pa[w][i]; pv != pw {104				v, w = pv, pw105			}106		}107		return pa[v][0]108	}109	isAncestor := func(f, v int) bool { return nodes[f].l < nodes[v].l && nodes[v].l <= nodes[f].r }110 111	t := make(fenwickDiff16, n+1)112	for ; q > 0; q-- {113		Fscan(in, &op, &v)114		v--115		if op == 1 {116			rt = v117		} else if op == 2 {118			Fscan(in, &w, &val)119			w--120			lca := getLCA(v, w)121			if lca == rt {122				t.add(1, n, val)123			} else if !isAncestor(lca, rt) {124				// 更新 lca 子树125				p := nodes[lca]126				t.add(p.l, p.r, val)127			} else { // lca 是 rt 的祖先128				maxD := max(dep[getLCA(rt, v)], dep[getLCA(rt, w)])129				if maxD < dep[rt] {130					subV := uptoDep(rt, maxD+1)131					p := nodes[subV]132					t.add(p.l, p.r, -val)133				}134				// 更新整棵树135				t.add(1, n, val)136			}137		} else {138			if v == rt {139				Fprintln(out, a[0]+t.query(1, n))140			} else if !isAncestor(v, rt) {141				p := nodes[v]142				Fprintln(out, a[v]+t.query(p.l, p.r))143			} else { // v 是 rt 的祖先144				subV := uptoDep(rt, dep[v]+1)145				p := nodes[subV]146				Fprintln(out, a[0]+t.query(1, n)-a[subV]-t.query(p.l, p.r))147			}148		}149	}150}151 152//func main() { cf916E(os.Stdin, os.Stdout) }153