【例题记录】树链剖分求树上区间询问

pre文章：树链剖分 - OI Wiki

配套习题：https://loj.ac/p/10138

#include<bits/stdc++.h>
using namespace std;
const int N = 60010;
const int INF = 0x3f3f3f3f;
int head[N], tot, w[N];
int siz[N], top[N], son[N], dep[N], fa[N], dfn[N], rnk[N], cnt;
struct node {
	int v, ne;
}e[N<<1];
struct treenode {
	int l, r, max, sum;
}tree[N<<2];
void add(int u, int v) {
	e[++tot].v = v, e[tot].ne = head[u], head[u] = tot;
}
void build(int root, int left, int right) {
	if (left == right) {
		tree[root] = { left,right,w[rnk[left]],w[rnk[left]] }; return;
	}
	tree[root] = { left,right,0,0 };
	int mid = left + right >> 1;
	build(root << 1, left, mid);
	build(root << 1 | 1, mid + 1, right);
	tree[root].max = max(tree[root << 1].max, tree[root << 1 | 1].max);
	tree[root].sum = tree[root << 1].sum + tree[root << 1 | 1].sum;
}
int querymax(int root, int l, int r) {
	if (l <= tree[root].l && tree[root].r <= r) return tree[root].max;
	int mid = tree[root].l + tree[root].r >> 1;
	if (r <= mid) return querymax(root << 1, l, r);
	else if (l > mid) return querymax(root << 1 | 1, l, r);
	else return max(querymax(root << 1, l, r), querymax(root << 1 | 1, l, r));
}
int querysum(int root, int l, int r) {
	if (l <= tree[root].l && tree[root].r <= r) return tree[root].sum;
	int mid = tree[root].l + tree[root].r >> 1;
	if (r <= mid) return querysum(root << 1, l, r);
	else if (l > mid) return querysum(root << 1 | 1, l, r);
	else return querysum(root << 1, l, r) + querysum(root << 1 | 1, l, r);
}
void update(int root, int dfn, int val) {
	if (tree[root].l == tree[root].r) {
		tree[root].max = val, tree[root].sum = val; return;
	}
	int mid = tree[root].l + tree[root].r >> 1;
	if (dfn <= mid) update(root << 1, dfn, val);
	else update(root << 1 | 1, dfn, val);
	tree[root].max = max(tree[root << 1].max, tree[root << 1 | 1].max);
	tree[root].sum = tree[root << 1].sum + tree[root << 1 | 1].sum;
}
void dfs1(int now) {
	siz[now] = 1;
	for (int i = head[now]; i; i = e[i].ne) {
		if (dep[e[i].v]) continue;
		dep[e[i].v] = dep[now] + 1;
		fa[e[i].v] = now;
		dfs1(e[i].v);
		siz[now] += siz[e[i].v];
		if (!son[now] || siz[son[now]] < siz[e[i].v]) son[now] = e[i].v;
	}
}
void dfs2(int now, int tp) {
	top[now] = tp;
	dfn[now] = ++cnt;
	rnk[cnt] = now;
	if (son[now]) dfs2(son[now], tp);
	for (int i = head[now]; i; i = e[i].ne) {
		if (e[i].v != son[now] && e[i].v != fa[now]) dfs2(e[i].v, e[i].v);
	}
}
int getmax(int x, int y) {
	int ret = -INF, fx = top[x], fy = top[y];
	while (fx != fy) {
		if (dep[fx] >= dep[fy]) {
			ret = max(ret, querymax(1, dfn[fx], dfn[x])), x = fa[fx];
		}
		else {
			ret = max(ret, querymax(1, dfn[fy], dfn[y])), y = fa[fy];
		}
		fx = top[x], fy = top[y];
	}
	if (dfn[x] < dfn[y]) ret = max(ret, querymax(1, dfn[x], dfn[y]));
	else ret = max(ret, querymax(1, dfn[y], dfn[x]));
	return ret;
}
int getsum(int x, int y) {
	int ret = 0, fx = top[x], fy = top[y];
	while (fx != fy) {
		if (dep[fx] >= dep[fy]) {
			ret += querysum(1, dfn[fx], dfn[x]), x = fa[fx];
		}
		else {
			ret += querysum(1, dfn[fy], dfn[y]), y = fa[fy];
		}
		fx = top[x], fy = top[y];
	}
	if (dfn[x] < dfn[y]) ret += querysum(1, dfn[x], dfn[y]);
	else ret += querysum(1, dfn[y], dfn[x]);
	return ret;
}
int main() {
	int n; cin >> n;
	for (int i = 1; i <= n - 1; i++) {
		int u, v; cin >> u >> v;
		add(u, v), add(v, u);
	}
	for (int i = 1; i <= n; i++) cin >> w[i];
	dep[1] = 1;
	dfs1(1);
	dfs2(1, 1);
	build(1, 1, n);
	int q; cin >> q;
	while (q--) {
		string op; int u, v; cin >> op >> u >> v;
		if (op == "CHANGE") update(1, dfn[u], v);
		else if (op == "QMAX") cout << getmax(u, v) << endl;
		else cout << getsum(u, v) << endl;
	}
	return 0;
}