树链剖分

洛谷p3384

code

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long LL;
#define lc p<<1
#define rc p<<1|1
const int N=1e5+10;
LL n,m,r,mod,a[N];
int dep[N],f[N];
//dfn[u],节点u在线段树的位置
//son重孩子,seg为线段树位置对应的节点
int dfn[N],son[N],siz[N],seg[N];
//top重链头
int cnt,top[N];
struct node{
    int l,r;
    LL sum;
    LL ad;
}tr[N<<2];
vector<int> edges[N];
void dfs1(int u,int fa){
    dep[u]=dep[fa]+1;
    f[u]=fa;
    for(int &v:edges[u]){
        if(v==fa)
            continue;
        dfs1(v,u);
    }
    siz[u]=1;
    int sid=0;
    for(int &v:edges[u]){
        if(v==fa)
            continue;
        siz[u]+=siz[v];
        if(sid==0||siz[v]>siz[sid]){
            sid=v;
        }
    }
    son[u]=sid;
}

void dfs2(int u,int t){
    top[u]=t;
    dfn[u]=++cnt;
    seg[cnt]=u;
    if(!son[u]){
        top[u]=t;
        return;
    }
    dfs2(son[u],t);
    for(int &v:edges[u]){
        if(v==f[u]||v==son[u])
            continue;
        dfs2(v,v);
    }
}

void pushup(int p){
    tr[p].sum=(tr[lc].sum+tr[rc].sum)%mod;
}

void build(int p,int l,int r){
    tr[p]={l,r,a[seg[l]]%mod};
    if(l==r)
        return;
    int mid=(l+r)>>1;
    build(lc,l,mid);
    build(rc,mid+1,r);
    pushup(p);
}

void lazy(int p,LL ad){
    int l=tr[p].l,r=tr[p].r;
    tr[p].sum=(tr[p].sum+ad*(r-l+1)%mod)%mod;
    (tr[p].ad+=ad)%=mod;
}

void pushdown(int p){
    if(tr[p].ad){
        lazy(lc,tr[p].ad);
        lazy(rc,tr[p].ad);
        tr[p].ad=0;
    }
}

LL query(int p,int x,int y){
    int l=tr[p].l,r=tr[p].r;
    if(x<=l&&r<=y){
        return tr[p].sum;
    }
    pushdown(p);
    int mid=(l+r)>>1;
    LL ret=0;
    if(y<=mid)
        ret=query(lc,x,y)%mod;
    else if(x>mid)
        ret=query(rc,x,y)%mod;
    else
        ret=(query(lc,x,y)+query(rc,x,y))%mod;
    pushup(p);
    return ret;
}

void modify(int p,int x,int y,LL ad){
    int l=tr[p].l,r=tr[p].r;
    if(x<=l&&r<=y){
        (tr[p].ad+=ad)%=mod;
        tr[p].sum=(tr[p].sum+(r-l+1)*ad%mod)%mod;
        return;
    }
    int mid=(l+r)>>1;
    pushdown(p);
    if(x<=mid)
        modify(lc,x,y,ad);
    if(y>mid)
        modify(rc,x,y,ad);
    pushup(p);
}
void addpath(LL x,LL y,LL z){
    while(top[x]!=top[y]){
        if(dep[top[x]]>dep[top[y]]){
            int fa=f[top[x]];
            modify(1,dfn[top[x]],dfn[x],z);
            x=fa;
        }
        else{
            int fa=f[top[y]];
            modify(1,dfn[top[y]],dfn[y],z);
            y=fa;
        }
    }

    if(dep[x]>dep[y])
        modify(1,dfn[y],dfn[x],z);
    else
        modify(1,dfn[x],dfn[y],z);
}

void addtree(LL x,LL z){
    int l=dfn[x],r=siz[x]+dfn[x]-1;
    modify(1,l,r,z);
}

LL pathsum(LL x,LL y){
    LL ret=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]>dep[top[y]]){
            int fa=f[top[x]];
            ret=(ret+query(1,dfn[top[x]],dfn[x]))%mod;
            x=fa;
        }
        else{
            int fa=f[top[y]];
            ret=(ret+query(1,dfn[top[y]],dfn[y]))%mod;
            y=fa;
        }
    }
    if(dep[x]>dep[y])
        ret=(ret+query(1,dfn[y],dfn[x]))%mod;
    else
        ret=(ret+query(1,dfn[x],dfn[y]))%mod;
    return ret;
}

LL treesum(LL x){
    int l=dfn[x],r=dfn[x]+siz[x]-1;
    return query(1,l,r);
}

int main(){
    cin.tie(nullptr)->sync_with_stdio(false);
    cin>>n>>m>>r>>mod;
    for(int i=1;i<=n;i++)
        cin>>a[i];
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        edges[u].push_back(v);
        edges[v].push_back(u);
    }
    dfs1(r,0);
    dfs2(r,r);
    build(1,1,n);
    while(m--){
        int op;
        cin>>op;
        if(op==1){
            LL x,y,z;
            cin>>x>>y>>z;
            addpath(x,y,z);
        }
        else if(op==2){
            LL x,y;
            cin>>x>>y;
            cout<<pathsum(x,y)<<'\n';
        }
        else if(op==3){
            LL x,z;
            cin>>x>>z;
            addtree(x,z);
        }
        else{
            LL x;
            cin>>x;
            cout<<treesum(x)<<'\n';
        }
    }
    return 0;
}