树链剖分模板

测试题目 P3384 【模板】树链剖分

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> piir;
const int maxn = 1e5+5;

int n,m,r,mod;
ll wt[maxn],w[maxn];//点权
struct node{
    int v,next;
}e[maxn*2];
int head[maxn],tot;
void init(){
    tot=0;
    memset(head,-1,sizeof(head));
}
void add(int u,int v){
    e[tot].v=v;
    e[tot].next=head[u];
    head[u]=tot++;
}

//segment
#define mid (l+r>>1)
#define ls  (k1<<1)
#define rs  (k1<<1|1)
#define len (r-l+1)
ll sum[maxn<<2],lazy[maxn<<2];

void pushup(int k1){sum[k1]=(sum[ls]+sum[rs])%mod;}
void pushdown(int k1,int l,int r){
    lazy[ls] = (lazy[ls]+lazy[k1]) % mod;
    lazy[rs] = (lazy[rs]+lazy[k1]) % mod;
    sum[ls]  = (sum[ls] + (lazy[k1]*(mid-l+1))%mod) % mod;
    sum[rs]  = (sum[rs] + (lazy[k1]*(r-mid))%mod) % mod;
    lazy[k1] = 0;
}
void build(int k1,int l,int r){
    if(l==r) {sum[k1]=wt[l];return ;}
    build(ls,l,mid);
    build(rs,mid+1,r);
    pushup(k1);
}
void update(int k1,int l,int r,int L,int R,ll v){
    if(L<=l && r<=R) {
        sum[k1] = (sum[k1]+(v*len)%mod) % mod;
        lazy[k1] = (lazy[k1]+v) % mod;
        return ;
    }
    if(lazy[k1]) pushdown(k1,l,r);
    if(L<=mid) update(ls,l,mid,L,R,v);
    if(R >mid) update(rs,mid+1,r,L,R,v);
    pushup(k1);
}
ll query(int k1,int l,int r,int L,int R){
    if(L<=l && r<=R) return sum[k1];
    if(lazy[k1]) pushdown(k1,l,r);
    ll res=0;
    if(L<=mid) res = (res+query(ls,l,mid,L,R)) % mod;
    if(R >mid) res = (res+query(rs,mid+1,r,L,R)) % mod;
    return res;
}

//hexvy
int sz[maxn],son[maxn],dep[maxn],fa[maxn];//子树大小，重儿子编号，深度，父亲节点编号
int id[maxn],top[maxn],cnt;               //新编号，当前链顶端节点，dfs_clock编号

void dfs1(int u,int f,int deep){
    dep[u]=deep;
    fa[u]=f;
    sz[u]=1;
    son[u]=0;
    for(int i=head[u];i!=-1;i=e[i].next){
        int v=e[i].v;
        if(v==f) continue;
        dfs1(v,u,deep+1);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int topf){
    id[u]=++cnt;
    wt[cnt]=w[u];
    top[u]=topf;
    if(!son[u]) return ; //每个节点都应该有一个重儿子，没有则为叶子节点
    dfs2(son[u],topf);   //先处理重儿子
    for(int i=head[u];i!=-1;i=e[i].next){
        int v=e[i].v;
        if(v==fa[u] || v==son[u]) continue;
        dfs2(v,v);       //每个轻儿子都有一条以itself开始的链
    }
}
ll qRange(int x,int y){
    ll ans=0;
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        ans += query(1,1,n,id[top[x]],id[x]);
        ans %= mod;
        x=fa[top[x]];
    }
    //在一条链上
    if(dep[x] > dep[y]) swap(x,y);
    ans += query(1,1,n,id[x],id[y]);
    ans %= mod;
    return ans;
}
ll qSon(int x){
    return query(1,1,n,id[x],id[x]+sz[x]-1);
}
void updRange(int x,int y,ll v){
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        update(1,1,n,id[top[x]],id[x],v);
        x=fa[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    update(1,1,n,id[x],id[y],v);
}
void updSon(int x,ll v){
    update(1,1,n,id[x],id[x]+sz[x]-1,v);
}

int main(){
    //freopen("in.txt","r",stdin);
    scanf("%d%d%d%lld",&n,&m,&r,&mod);
    init();
    for(int i=1;i<=n;i++){
        scanf("%lld",&w[i]);
        w[i]%=mod;
    }
    for(int u,v,i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }

    dfs1(r,0,1);
    dfs2(r,r);
    build(1,1,n);

    for(int i=1;i<=m;i++){
        int op,x,y;
        ll z;
        scanf("%d",&op);
        if(op==1){
            scanf("%d%d%lld",&x,&y,&z);
            updRange(x,y,z%mod);
        }
        else if(op==2){
            scanf("%d%d",&x,&y);
            printf("%lld\n",qRange(x,y));
        }
        else if(op==3){
            scanf("%d%lld",&x,&z);
            updSon(x,z%mod);
        }
        else if(op==4){
            scanf("%d",&x);
            printf("%lld\n",qSon(x));
        }
    }
    return 0;
}