树链剖分（以维护线段树为例）

Posted 2020-10-06

关于树链剖分的有关知识：http://www.cnblogs.com/sagitta/p/5660749.html

以下是洛谷p3384经过无数次WA和RE（最后发现只是有一个函数的调用写反了qwq）终于AC的代码：

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cctype>
#include<climits>
#include<ctime>
#include<queue>
#include<vector>
#include<map>
#include<algorithm>
#include<iomanip>
using namespace std;
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define dep(i,a,b) for(int i=a;i>=b;i--)
typedef long long LL;
inline int read(){
    int x=0;char ch=getchar();
    while(ch<‘0‘||ch>‘9‘)ch=getchar();
    while(ch>=‘0‘&&ch<=‘9‘){
        x=x*10+ch-‘0‘;
        ch=getchar();
    }
    return x;
}
const int M=100001;
LL sum[M<<2],add[M<<2];
int n,m,root,p,cnt=1,tot=0,a[M],son[M],dep[M],top[M],siz[M],dfn[M],dfx[M],fa[M],to[M<<1],head[M],next[M<<1];
/*
son[u]为u的重儿子,dep[u]为u的深度,top[u]为u所在重链的根节点,siz[u]为以u为根的子树的节点个数,dfn[u]为u的dfs序,dfx[x]为dfs序等于x的节点,fa[u]为u的父亲节点,to,head,next数组维护链式前向星
*/

//线段树
void getup(int rt){
    sum[rt]=sum[rt<<1]+sum[rt<<1|1];
}
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
void clean(int rt,int len){
    if(add[rt]){
        sum[rt<<1]+=add[rt]*(len-(len>>1));
        sum[rt<<1|1]+=add[rt]*(len>>1);
        add[rt<<1]+=add[rt];
        add[rt<<1|1]+=add[rt];
        add[rt]=0;
    }
}
void build(int l,int r,int rt){
    add[rt]=0;
    if(l==r){
        sum[rt]=a[dfx[l]];//按照节点的dfs序建立线段树
        return;
    }
    int m=(l+r)>>1;
    build(lson);
    build(rson);
    getup(rt);
}
void update(int l,int r,int rt,int L,int R,LL k){
    int len=r-l+1;
    if(L<=l&&r<=R){
        sum[rt]+=(LL)len*k;
        add[rt]+=k;
        return;
    }
    clean(rt,len);
    int m=(l+r)>>1;
    if(L<=m)update(lson,L,R,k);
    if(m<R)update(rson,L,R,k);
    getup(rt);
}
LL query(int l,int r,int rt,int L,int R){
    if(L<=l&&r<=R)return sum[rt];
    clean(rt,r-l+1);
    int m=(l+r)>>1;
    LL ans=0;
    if(L<=m)ans+=query(lson,L,R);
    if(m<R)ans+=query(rson,L,R);
    return ans;
}

//插入边
void Insert(int u,int v){
    to[++cnt]=v;
    next[cnt]=head[u];
    head[u]=cnt;
    to[++cnt]=u;
    next[cnt]=head[v];
    head[v]=cnt;
}

//第一次dfs找出每个节点的重儿子(叶子节点无重儿子)
void dfs1(int u){
    siz[u]=1;
    for(int i=head[u];i!=-1;i=next[i]){
        int v=to[i];
        if(v!=fa[u]){
            fa[v]=u;
            dep[v]=dep[u]+1;
            dfs1(v);
            siz[u]+=siz[v];
            if(!son[u]||siz[son[u]]<siz[v])son[u]=v;
        }
    }
}
//第二次dfs连重边成重链
void dfs2(int u){
    dfn[u]=++tot;
    dfx[tot]=u;
    if(son[u]){
        top[son[u]]=top[u];
        dfs2(son[u]);
    }
    for(int i=head[u];i!=-1;i=next[i]){
        int v=to[i];
        if(v!=son[u]&&v!=fa[u]){
            top[v]=v;
            dfs2(v);
        }
    }
}

//将节点x与y的最短路径上所有点权值都加上z
void Plus(int x,int y,int z){
    while(1){
        int tx=top[x],ty=top[y];
        if(tx==ty){//若x与y在同一条重链上,直接更改线段树信息
            if(dfn[x]<dfn[y])update(1,n,1,dfn[x],dfn[y],z);
            else update(1,n,1,dfn[y],dfn[x],z);
            return;
        }
    //不在一条重链上,则先处理所在重链的根节点深度较大的点,然后跳转到该根节点的父亲节点
        if(dep[tx]<dep[ty]){
            update(1,n,1,dfn[ty],dfn[y],z);
            y=fa[ty];
        }
        else{
            update(1,n,1,dfn[tx],dfn[x],z);
            x=fa[tx];
        }
    }
}

//询问节点x与y最短路径上所有节点的权值和,操作同上
LL Find(int x,int y){
    LL ans=0;
    while(1){
        int tx=top[x],ty=top[y];
        if(tx==ty){
            if(dfn[x]<dfn[y])ans+=query(1,n,1,dfn[x],dfn[y]);
            else ans+=query(1,n,1,dfn[y],dfn[x]);
            return ans;
        }
        if(dep[tx]<dep[ty]){
            ans+=query(1,n,1,dfn[ty],dfn[y]);
            y=fa[ty];
        }
        else{
            ans+=query(1,n,1,dfn[tx],dfn[x]);
            x=fa[tx];
        }
    }
}
int main(){
    n=read();m=read();root=read();p=read();
    memset(head,-1,sizeof(head));
    memset(son,0,sizeof(son));
    memset(fa,0,sizeof(fa));
    rep(i,1,n)a[i]=read();
    rep(i,1,n-1){
        int x=read(),y=read();
        Insert(x,y);
    }
    dep[root]=1;
    top[root]=root;
    dfs1(root);
    dfs2(root);
    build(1,n,1);
    while(m--){
        int q=read();
        if(q==1){
            int x=read(),y=read(),z=read();
            Plus(x,y,z);
        }
        else if(q==2){
            int x=read(),y=read();
            cout<<Find(x,y)%p<<endl;
        }
        else if(q==3){//将以x为根的子树上所有点权值+z
            int x=read(),z=read();
            update(1,n,1,dfn[x],dfn[x]+siz[x]-1,z);//以一个点为根的子树的所有点的dfs序是连续的,直接更改即可
        }
        else if(q==4){//询问以节点x为根的子树上所有点的权值和,同上
            int x=read();
            cout<<query(1,n,1,dfn[x],dfn[x]+siz[x]-1)%p<<endl;
        }
    }
    return 0;
}