树链剖分
核心思路为将整棵树按照重儿子优先遍历形成DFS序,这样u-v的路径可以变为连续至多log(n)个连续区间,用线段树来维护区间的和即可。
功能:
- modify_path(u, v, k);将路径u->v的所有节点权值增加K
- modify_tree(u, k);//将k为根的子树所有节点增加K
- printf(“%lld\n”, ask_path(u, v));询问u->v的所有节点的和
- printf(“%lld\n”, ask_tree(u));寻找u为子树所有节点的和
#include <bits/stdc++.h>
using namespace std;
//#define map unordered_map
//#define int long long
const int N = 1e5 + 10, M = N << 1;
typedef long long ll;
int n, m;
int h[N], e[M], ne[M], w[N], idx;
int id[N]; //原来树中每个点在dfs序中的编号
int nw[N]; //每个编号点的权值,即新编号的点的权值 dfs序中第i个点编号
int cnt;
int depth[N]; //每个点所在深度
int sz[N]; //以每个点为根节点的子树大小
int top[N]; //每个点所在重链的顶点
int fa[N]; //每个点父节点
int son[N]; //每个点的重儿子
struct node
{
int l, r;
ll add, sum; //线段树中维护两个值
} t[N << 2];
void add(int a, int b) {
e[idx] = b, ne[idx] = h[a], h[a] = idx++;
}
void dfs1(int u, int father, int dep) //节点编号 其父节点编号 当前深度
{
depth[u] = dep, fa[u] = father, sz[u] = 1;
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
if (j == father) continue;
dfs1(j, u, dep + 1);
sz[u] += sz[j]; //当前子树大小加上第 j 棵子树大小
if (sz[son[u]] < sz[j]) son[u] = j; //如果当前重儿子节点个数小于第 j 棵子树的节点数,说明当前重儿子应当为 j
}
}
void dfs2(int u, int t) //当前点 以及当前点所在重链的顶点
{
id[u] = ++cnt; //dfs序
nw[cnt] = w[u]; //dfs序中第cnt个点权为w[u]
top[u] = t; //当前点所在重链的顶点是t
if (!son[u]) return; //如果当前点是叶子节点,即没有儿子
dfs2(son[u], t); //否则优先搜索其重儿子
//之后dfs所有轻儿子
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
//如果xx或j为其重儿子,由于重儿子已被搜过,那么跳过当前循环
if (j == fa[u] || j == son[u]) continue;
dfs2(j, j); //递归轻儿子,轻儿子所在重链顶点就是自己
}
}
void pushup(int u)
{
t[u].sum = t[u << 1].sum + t[u << 1 | 1].sum;
}
void pushdown(int u) //下传懒标记
{
auto &rt = t[u], &le = t[u << 1], &ri = t[u << 1 | 1];
if (rt.add)
{
le.add += rt.add, le.sum += rt.add * (le.r - le.l + 1);
ri.add += rt.add, ri.sum += rt.add * (ri.r - ri.l + 1);
rt.add = 0;
}
}
void build(int u, int l, int r)
{
t[u] = { l, r };
if (l == r) {
t[u].sum = nw[l];
return;
}
int mid = l + r >> 1;
build(u << 1, l, mid), build(u << 1 | 1, mid + 1, r);
pushup(u);
}
void modify(int u, int l, int r, int v)
{
if (l <= t[u].l && r >= t[u].r)
{
t[u].add += v, t[u].sum += v * (t[u].r - t[u].l + 1);
return;
}
pushdown(u);
int mid = t[u].l + t[u].r >> 1;
if (l <= mid) modify(u << 1, l, r, v);
if (r > mid) modify(u << 1 | 1, l, r, v);
pushup(u);
}
ll ask(int u, int l, int r)
{
if (l <= t[u].l && r >= t[u].r)
{
return t[u].sum;
}
pushdown(u);
int mid = t[u].l + t[u].r >> 1;
ll res = 0;
if (l <= mid) res += ask(u << 1, l, r);
if (r > mid) res += ask(u << 1 | 1, l, r);
return res;
}
void modify_path(int u, int v, int k) //爬山法
{
//如何判断两点是否在同一条重链中?
//类似于并查集 存了每个点所在重链的顶点编号
//判断两点所在重链顶点是否一样即可
while (top[u] != top[v]) //当两点不在同一条重链中时
{
if (depth[top[u]] < depth[top[v]]) swap(u, v);
//优先走u所在重链
modify(1, id[top[u]], id[u], k); //修改这段连续区间 即子树
u = fa[top[u]]; //跳到重链上方
}
if (depth[u] < depth[v]) swap(u, v);
modify(1, id[v], id[u], k); //修改最后一段
}
ll ask_path(int u, int v) //与修改形式一样
{
ll res = 0;
while (top[u] != top[v])
{
if (depth[top[u]] < depth[top[v]]) swap(u, v);
res += ask(1, id[top[u]], id[u]);
u = fa[top[u]];
}
if (depth[u] < depth[v]) swap(u, v);
res += ask(1, id[v], id[u]);
return res;
}
void modify_tree(int u, int v) //以u为子树是一段连续区间,左右端点如下
{
modify(1, id[u], id[u] + sz[u] - 1, v);
}
ll ask_tree(int u)
{
return ask(1, id[u], id[u] + sz[u] - 1);
}
signed main()
{
cin >> n;
for (int i = 1; i <= n; ++i)
{
scanf("%d", &w[i]);
}
memset(h, -1, sizeof h);
int t = n - 1;
while (t--)
{
int u, v; scanf("%d%d", &u, &v);
add(u, v), add(v, u);
}
dfs1(1, -1, 1); //先求每个点的重儿子
dfs2(1, -1); //求一下 dfs 序
build(1, 1, n); //建线段树
//树链剖分
cin >> m;
while (m--)
{
int t, u, v, k; scanf("%d%d", &t, &u);
if (t == 1)
{
scanf("%d%d", &v, &k);
modify_path(u, v, k);//将路径u->v的所有节点权值增加K
}
else if (t == 2)
{
scanf("%d", &k);
modify_tree(u, k);//将k为根的子树所有节点增加K
}
else if (t == 3)
{
scanf("%d", &v);
printf("%lld\n", ask_path(u, v));//询问u->v的所有节点的和
}
else
{
printf("%lld\n", ask_tree(u));//寻找u为子树所有节点的和
}
}
return 0;
}