树链剖分模板

摘自博客：https://blog.csdn.net/VictoryCzt/article/details/85268376

博主讲得十分的清楚

以下是带有注释的代码：

/*
这是一份自用的树链剖分模板，
自己加的注释，便于更好的理解
该模板支持的操作：
1.换根
2.一条链上点权加
3.一个子树内点权加
4.询问一条链上点权的和
5.询问一个子树内的点权和
6.求LCA
*/
#include<bits/stdc++.h>
#define numm ch-48
#define pd putchar(' ')
#define pn putchar('
')
#define pb push_back
#define debug(args...) cout<<#args<<"->"<<args<<endl
#define bug cout<<"************"
using namespace std;
template <typename T>
void read(T &res) {
    bool flag=false;char ch;
    while(!isdigit(ch=getchar())) (ch=='-')&&(flag=true);
    for(res=numm;isdigit(ch=getchar());res=(res<<1)+(res<<3)+numm);
    flag&&(res=-res);
}
template <typename T>
void write(T x) {
    if(x<0) putchar('-'),x=-x;
    if(x>9) write(x/10);
    putchar(x%10+'0');
}
typedef long long ll;
const int inf=0x3f3f3f3f;
const int maxn=1e5+10;
/*建图部分*/
struct node {///链式前向星
    int v,net;
    node(){}
    node(int v,int net):v(v),net(net){}
}e[maxn<<1];///树有双向边，边数*2
int head[maxn],cnt,n;
void init() {
    cnt=0;
    for(int i=1;i<=n;i++)
        head[i]=-1;
}
void addedge(int u,int v) {///建双向边
    e[cnt]=node(v,head[u]);
    head[u]=cnt++;
    e[cnt]=node(u,head[v]);
    head[v]=cnt++;
}
int id[maxn],rk[maxn],fa[maxn],top[maxn],dep[maxn],son[maxn],sze[maxn],val[maxn],root=1;
//id[u]:点u的dfs序，一条重链上的点标号连续
//sze[u]:点u的子树结点个数，包含点u
//son[u]:点u的重儿子的id
//fa[u]:点u的父亲,根结点1的父亲为0
//top[u]:点u所在重链的头结点,轻链则为本身，则top[u]=u;
//val[u]:点u的初始权值
//rk[id[u]]:id[u]对应的u结点
//dep[u]:dfs时，u对应的深度
int minp[maxn<<2];
ll sum[maxn<<2],lazy[maxn<<2];
/*线段树部分*/
//sum[k]:点k所在区间的树结点总和
//lazy[k]:点k所在区间的树结点总和的懒标记
//minp[k]:点k所在区间内的最小深度的点u的下标
int get_minp(int a,int b) { ///返回a,b结点中深度最小的结点
    if(a==inf) return b;
    if(b==inf) return a;
    if(dep[a]<dep[b]) return a;
    return b;
}
void pushup(int k) {
    sum[k]=sum[k<<1]+sum[k<<1|1];
    minp[k]=get_minp(minp[k<<1],minp[k<<1|1]);
}
void pushdown(int k,int l,int r,int mid) {
    if(!lazy[k]) return ;
    sum[k<<1]+=(mid-l+1)*lazy[k];
    sum[k<<1|1]+=(r-mid)*lazy[k];
    lazy[k<<1]+=lazy[k];
    lazy[k<<1|1]+=lazy[k];
    lazy[k]=0;
}
void build(int k,int l,int r) {///建树
    if(l==r) {
        sum[k]=val[rk[l]];
        minp[k]=rk[l];
        return ;
    }
    int mid=l+r>>1;
    build(k<<1,l,mid);
    build(k<<1|1,mid+1,r);
    pushup(k);
}
void update(int k,int l,int r,int ql,int qr,ll w) {///区间修改
    if(ql<=l&&r<=qr) {
        sum[k]+=(r-l+1)*w;
        lazy[k]+=w;
        return ;
    }
    int mid=l+r>>1;
    pushdown(k,l,r,mid);///attention!!!
    if(ql<=mid) update(k<<1,l,mid,ql,qr,w);
    if(mid<qr) update(k<<1|1,mid+1,r,ql,qr,w);
    pushup(k);
}
int query_pos(int k,int l,int r,int ql,int qr) {///区间查询最小深度的结点
    if(ql>qr) return inf;
    if(ql<=l&&r<=qr) return minp[k];
    int mid=l+r>>1;
    if(qr<=mid) return query_pos(k<<1,l,mid,ql,qr);
    else if(ql>mid) return query_pos(k<<1|1,mid+1,r,ql,qr);
    return get_minp(query_pos(k<<1,l,mid,ql,qr),query_pos(k<<1|1,mid+1,r,ql,qr));
}
ll query_sum(int k,int l,int r,int ql,int qr) {///区间求和
    if(ql<=l&&r<=qr) return sum[k];
    int mid=l+r>>1;
    pushdown(k,l,r,mid);
    if(qr<=mid) return query_sum(k<<1,l,mid,ql,qr);
    else if(ql>mid) return query_sum(k<<1|1,mid+1,r,ql,qr);
    else return query_sum(k<<1,l,mid,ql,qr)+query_sum(k<<1|1,mid+1,r,ql,qr);
}
/*树剖部分*/
void dfs1(int u) {  ///得到每个点的深度,父亲,，子树大小和重儿子
    sze[u]=1;   ///把自己算进去
    for(int i=head[u];~i;i=e[i].net) {
        int v=e[i].v;
        if(v==fa[u]) continue;
        fa[v]=u;
        dep[v]=dep[u]+1;
        dfs1(v);
        sze[u]+=sze[v]; ///回溯累加
        if(!son[u]||sze[son[u]]<sze[v])///更新重儿子
            son[u]=v;
    }
}
void dfs2(int u,int root) {
    top[u]=root;
    id[u]=++cnt;
    rk[cnt]=u;
    if(!son[u]) return ;
    dfs2(son[u],root);
    ///先dfs重儿子，使得每条重链上的结点标号连续
    for(int i=head[u];~i;i=e[i].net) {
        int v=e[i].v;
        if(v==son[u]||v==fa[u]) continue;
        dfs2(v,v);///轻链的top则为本身
    }
}
int lca(int a,int b) {  ///求公共祖先
    while(top[a]!=top[b]) {
        if(dep[top[a]]<dep[top[b]]) swap(a,b);
        a=fa[top[a]];
    }
    if(dep[a]>dep[b]) swap(a,b);
    return a;
}
int Upto(int a,int b) { ///返回a到b上深度最小的点的下标
    int ans=inf;
    while(top[a]!=top[b]) {
        if(dep[top[a]]<dep[top[b]]) swap(a,b);
        ///不能包含u点,要分类讨论
        if(top[a]!=b) ans=get_minp(ans,query_pos(1,1,n,id[top[a]],id[a]));
        else ans=get_minp(ans,query_pos(1,1,n,id[top[a]]+1,id[a]));
        a=fa[top[a]];
    }
    if(dep[a]>dep[b]) swap(a,b);
    
    ans=get_minp(ans,query_pos(1,1,n,id[a]+1,id[b]));///不能包含u点,所以+1
    return ans;
}
int check(int u) {  ///判r与root的关系
    if(u==root) return 0;///此时操作范围就是整个树，将整个树加或者求和即可。
    int L=lca(u,root);
    if(L==u) return Upto(u,root);/*当r在u的子树内时，操作范围就为
    整个树减去u到r路径上深度最小的点（不包含u点）的原来的子树*/
    else return -1;///当r不在u的原来的子树里面时，操作范围就是原来的子树。
}
void add_chain(int a,int b,ll w) { ///一条链上点权加
    while(top[a]!=top[b]) {
        if(dep[top[a]]<dep[top[b]]) swap(a,b);///找到深度最大的点
        update(1,1,n,id[top[a]],id[a],w);
        a=fa[top[a]];   ///一次跳一个
    }
    if(dep[a]>dep[b]) swap(a,b);
    update(1,1,n,id[a],id[b],w);
}
ll ask_chain(int a,int b) { ///询问一条链上点权的和
    ll ans=0;
    while(top[a]!=top[b]) {
        if(dep[top[a]]<dep[top[b]]) swap(a,b);
        ans+=query_sum(1,1,n,id[top[a]],id[a]);
        a=fa[top[a]];
    }
    if(dep[a]>dep[b]) swap(a,b);
    ans+=query_sum(1,1,n,id[a],id[b]);
    return ans;
}
void add_tree(int u,ll w) { ///一个子树内点权加
    int type=check(u);  ///判断u与root的情况
    if(!type) update(1,1,n,1,n,w);
    else if(type>0) {
        update(1,1,n,1,n,w);
        update(1,1,n,id[type],id[type]+sze[type]-1,-w);
    }
    else update(1,1,n,id[u],id[u]+sze[u]-1,w);
}
ll ask_tree(int u) {    ///询问u的子树内的点权和
    int type=check(u); ///判断u与root的情况
    if(!type) return query_sum(1,1,n,1,n);
    else if(type>0){
        ll ans=query_sum(1,1,n,id[type],id[type]+sze[type]-1);
        return query_sum(1,1,n,1,n)-ans;
    }
    else return query_sum(1,1,n,id[u],id[u]+sze[u]-1);
}
int main()
{
    int a,b;
    ll w;
    read(n);
    init();
    for(int i=1;i<=n;i++) ///得到每个点的初始权值
        read(val[i]);
    for(int i=2;i<=n;i++) { ///得到每个点的父亲
        read(a);
        addedge(a,i);   ///建边
    }
    dfs1(root);
    cnt=0;  ///attention!!!
    dfs2(1,1);
    build(1,1,n);
    int m;
    read(m);
    while(m--) {
        int oper;
        read(oper);
        if(oper==1) {   ///换根
            read(a);
            root=a;
        }
        else if(oper==2) {  ///一条链上点权加
            read(a),read(b),read(w);
            add_chain(a,b,w);
        }
        else if(oper==3) {   ///一个子树内点权加
            read(a),read(w);
            add_tree(a,w);
        }
        else if(oper==4) {  ///询问一条链上点权的和
            read(a),read(b);
            write(ask_chain(a,b));pn;
        }
        else if(oper==5) { ///询问一个子树内的点权和
            read(a);
            write(ask_tree(a));pn;
        }
    }
    return 0;
}