树链剖分模板

树链剖分

核心思路为将整棵树按照重儿子优先遍历形成DFS序，这样u-v的路径可以变为连续至多log（n）个连续区间，用线段树来维护区间的和即可。
功能：

1. modify_path(u, v, k);将路径u->v的所有节点权值增加K
2. modify_tree(u, k);//将k为根的子树所有节点增加K
3. printf(“%lld\n”, ask_path(u, v));询问u->v的所有节点的和
4. 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;
}