CF916E 换根树上问题

Link

题意：对一棵树进行三种操作。

• 把根设为 $x$。
• 将以 $lca(y,z)$ 为根的子树中所有点的权值加 $v$。
• 查询以 $x$ 为根的子树点权之和。

初始根为 $1$。

设相对于当前 $root$ 的所要操作（查询）子树的根节点为 $U$。

这里用 $[in[x],out[x]]$ 表示以 $x$ 为根的子树在 $dfs$ 序上的出现区间。

情况一：$U=root$。#

修改整棵树。

情况二：直接在原树上修改子树 $U$。#

满足以下条件之一：

• 整棵 $U$ 在到 $root$ 前已经遍历完了，$out[U]<in[root]$。
• 整棵 $root$ 在到 $U$ 前已经遍历完了，$in[U]>out[root]$。
• $U$ 是 $root$ 的后代，$in[root]<in[U]<out[U]<out[root]]$。

二三可以合写为 $in[U]>in[root]$。

情况三：$root$ 为 $U$ 的后代。#

先整树修改，在扣除 $U$ 在 $root$ 方向的子树的影响。
如何求这个方向子树的根，可以借鉴 $lca$ 的思想。

int get(int x, int y) {
	if(dep[x] < dep[y]) swap(x, y);
	per(i, 17, 0) if(dep[fa[x][i]] > dep[y]) x = fa[x][i];
	return x;
}

操作三的 $U$ 就是 $x$。
对于操作二，$U$ 为 $lca(y,z)$，$lca(root,y)$，$lca(root,z$ 中深度最大的点，可以自己模拟一下。

int lca(int a, int b, int c) {
	int x = lca(a, b), y = lca(a, c);
	if(dep[x] < dep[y]) x = y;
	y = lca(b, c);
	if(dep[x] < dep[y]) x = y;
	return x;
}

具体做法就是线段树维护 $dfs$ 序 + 大分讨，细节较多。

完整代码

#include<bits/stdc++.h>
#define rep(i, a, b) for(int i = (a); i <= (b); ++ i)
#define per(i, a, b) for(int i = (a); i >= (b); -- i)
#define pb emplace_back
#define All(X) X.begin(), X.end()
using namespace std;
using ll = long long;
constexpr int N = 1e5 + 5;

int n, m, dep[N], in[N], out[N], dfn;
int fa[N][18];
ll a[N];
vector<int> G[N];

struct Node {
	ll sum, tag;
	int sz;
} t[N << 2];

#define ls x << 1
#define rs ls | 1

void build(int x = 1, int l = 1, int r = n) {
	t[x].sz = r - l + 1;
	if(l == r) return;
	int mid = l + r >> 1;
	build(ls, l, mid), build(rs, mid + 1, r);
}

void pushup(int x) {
	t[x].sum = t[ls].sum + t[rs].sum;
}

void pushdown(int x) {
	if(t[x].tag) {
		t[ls].tag += t[x].tag, t[ls].sum += t[x].tag * t[ls].sz;
		t[rs].tag += t[x].tag, t[rs].sum += t[x].tag * t[rs].sz;
		t[x].tag = 0;
	}
}

void add(int L, int R, ll v, int x = 1, int l = 1, int r = n) {
	if(L > R) return;
	if(L <= l && r <= R) {
		t[x].sum += v * t[x].sz;
		t[x].tag += v;
		return;
	}
	pushdown(x);
	int mid = l + r >> 1;
	if(mid >= L) add(L, R, v, ls, l, mid);
	if(mid < R) add(L, R, v, rs, mid + 1, r);
	pushup(x);
}

ll query(int L, int R, int x = 1, int l = 1, int r = n) {
	if(L > R) return 0;
	if(L <= l && r <= R) {
		return t[x].sum;
	}
	pushdown(x);
	int mid = l + r >> 1;
	ll ret = 0;
	if(mid >= L) ret += query(L, R, ls, l, mid);
	if(mid < R) ret += query(L, R, rs, mid + 1, r);
	return ret;
}

void dfs(int x) {
	dep[x] = dep[fa[x][0]] + 1;
	in[x] = ++ dfn;
	add(dfn, dfn, a[x]);
	
	for(auto y : G[x]) {
		if(y != fa[x][0]) {
			fa[y][0] = x;
			rep(i, 1, 17) fa[y][i] = fa[fa[y][i - 1]][i - 1];
			dfs(y);
		}
	}
	out[x] = dfn;
}

int lca(int x, int y) {
	if(dep[x] < dep[y]) swap(x, y);
	per(i, 17, 0) if(dep[fa[x][i]] >= dep[y] ) x = fa[x][i];
	if(x == y) return x;
	per(i, 17, 0) if(fa[x][i] != fa[y][i]) x = fa[x][i], y = fa[y][i];
	return fa[x][0];
}
int lca(int a, int b, int c) {
	int x = lca(a, b), y = lca(a, c);
	if(dep[x] < dep[y]) x = y;
	y = lca(b, c);
	if(dep[x] < dep[y]) x = y;
	return x;
}
int get(int x, int y) {
	if(dep[x] < dep[y]) swap(x, y);
	per(i, 17, 0) if(dep[fa[x][i]] > dep[y]) x = fa[x][i];
	return x;
}

int root = 1;

int main() {
	ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
	cin >> n >> m;
	rep(i, 1, n) cin >> a[i];
	rep(i, 2, n) {
		int x, y; cin >> x >> y;
		G[x].pb(y);
		G[y].pb(x);
	}
	build(), dfs(1);
	rep(i, 1, m) {
		int op; cin >> op;
		if(op == 1) {
			cin >> root;
		}
		if(op == 2) {
			int y, z, v; cin >> y >> z >> v;
			int x = lca(root, y, z);
			
			if(x == root) {
				add(1, n, v);
			}
			else {
				if(in[x] > in[root] || out[x] < in[root]) {
					add(in[x], out[x], v);
				}
				else {
					int Fa = get(x, root);
					add(1, n, v);
					add(in[Fa], out[Fa], -v);
				}
			}
		}
		if(op == 3) {
			int x; cin >> x;
			if(x == root) {
				cout << t[1].sum << '\n';
			}
			else {
				if(in[x] > in[root] || out[x] < in[root]) {
					cout << query(in[x], out[x]) << '\n';
				}
				else {
					int Fa = get(x, root);
					cout << t[1].sum - query(in[Fa], out[Fa]) << '\n';
				}
			}
		}		
	}
	return 0;
}