P3384 【模板】重链剖分

树链剖分学习

#include <bits/stdc++.h>
using namespace std;
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
const int N=1e6+10;

struct egde{
    int to,next;
}e[N];
int head[N],tot;
void add(int u,int v){
    e[++tot]={v,head[u]};
    head[u]=tot;
}

int w[N];
int n,m,r,mod;

int wt[N];//树链剖分后新序号的权值 在线段树中用到
/*-----------------线段树---------------------*/
int a[N];
int lazy[N];
void pushup(int rt){
    a[rt]=(a[rt<<1]+a[rt<<1|1])%mod;
}

void build(int l,int r,int rt){
    if(l==r){
        a[rt]=wt[l]%mod;
        return;
    }
    int mid=l+r>>1;
    build(lson);
    build(rson);
    pushup(rt);
}

void pushdown(int rt,int len){
    if(lazy[rt]){
        lazy[rt<<1]+=lazy[rt];
        lazy[rt<<1]%=mod;

        lazy[rt<<1|1]+=lazy[rt];
        lazy[rt<<1|1]%=mod;

        a[rt<<1]+=lazy[rt]*(len-(len>>1))%mod;
        a[rt<<1|1]+=lazy[rt]*(len>>1)%mod;
        a[rt<<1]%=mod;
        a[rt<<1|1]%=mod;

        lazy[rt]=0;
    }
}
void update(int L,int R,int val,int l,int r,int rt){
    if(L<=l && r<=R){
        lazy[rt]+=val;
        a[rt]+=val*(r-l+1)%mod;
        return;
    }
    pushdown(rt,r-l+1);
    int mid=l+r>>1;
    if(L<=mid) update(L,R,val,lson);
    if(R>mid) update(L,R,val,rson);
    pushup(rt);
}

int query(int L,int R,int l,int r,int rt){
    if(L<=l && r<=R){
        return a[rt]%mod;
    }
    pushdown(rt,r-l+1);
    int mid=l+r>>1,res=0;
    if(L<=mid) res=(res+query(L,R,lson))%mod;
    if(R>mid) res=(res+query(L,R,rson))%mod;
    return res%mod;
}

/*----------------树链剖分----------------------*/
int dep[N],son[N],f[N],sz[N];//深度 重儿子 父节点 子树大小
void dfs1(int u,int fa){
    dep[u]=dep[fa]+1;
    f[u]=fa;
    sz[u]=1;//大小包括本身
    for(int i=head[u];i;i=e[i].next){
        int v=e[i].to;
        if(v!=fa){
            dfs1(v,u);
            sz[u]+=sz[v];
            if(sz[v]>sz[son[u]]) son[u]=v;
        }
    }
}

int dfn;//时间戳
int top[N],id[N];//每个链的链顶 树链剖分dfs2后新的编号
void dfs2(int u,int t){//当前节点 链最顶端的节点
    id[u]=++dfn;
    wt[dfn]=w[u];//把原来的点的权值赋值给新的编号
    top[u]=t;//记录u所在的链的链顶
    if(son[u]) dfs2(son[u],t);//如果有重儿子先走重儿子
    for(int i=head[u];i;i=e[i].next){
        int v=e[i].to;
        if(v!=f[u] && v!=son[u]){//v不是父节点 也不是重儿子 
            dfs2(v,v);//轻链的重链顶是自己
        }
    }
}

void updRange(int x,int y,int val){//(x,y)路径上所有点权都加val
    val%=mod;
    while(top[x]!=top[y]){//当两个点不在同一条链上
        if(dep[top[x]] < dep[top[y]])
            swap(x,y);//默认x的链顶比y的链顶深
        update(id[top[x]],id[x],val,1,n,1);//链顶到x的序是连续的
        x=f[top[x]];//往链顶的父节点更新
    }
    if(dep[x]>dep[y])//直到两点处于一条链上
        swap(x,y);//把深度更深的点作为y点
    update(id[x],id[y],val,1,n,1);
}

int qRange(int x,int y){//返回(x,y)路径上点权和
    int res=0;
    while(top[x]!=top[y]){//当两个点不在同一条链上
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        res+=query(id[top[x]],id[x],1,n,1);
        res%=mod;
        x=f[top[x]];
    }
    if(dep[x]>dep[y])swap(x,y);
    res+=query(id[x],id[y],1,n,1);
    return res%mod;
}

void updSon(int x,int k){
    update(id[x],id[x]+sz[x]-1,k,1,n,1);//子树区间右端点为id[x]+sz[x]-1
}

int qSon(int x){
    return query(id[x],id[x]+sz[x]-1,1,n,1);
}

int main(){
    ios::sync_with_stdio(0);

    cin>>n>>m>>r>>mod;//r是root
    for (int i = 1; i <= n; ++i) {
        cin>>w[i];
    }
    for (int i = 1,u,v; i < n; ++i) {
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
	//树链剖分
    dfs1(r,0);
    dfs2(r,r);
    //建立线段树
    build(1,n,1);
    
    for (int i = 1,k,x,y,z; i <= m; ++i) {
        cin>>k;
        if(k==1){
            cin>>x>>y>>z;
            updRange(x,y,z);//(x,y)路径上的点权都加上z
        }else if(k==2){
            cin>>x>>y;
            cout <<qRange(x,y) << endl;//问(x,y)路径上的点权和
        }else if(k==3){
            cin>>x>>y;
            updSon(x,y);//给x和它的子树都加上y
        }else{
            cin>>x;
            cout <<qSon(x) << endl;//问x和其子树的点权和
        }
    }
    return 0;
}