CF916E 换根树上问题
题意:对一棵树进行三种操作。
- 把根设为 \(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;
}