树链剖分(轻重树)模板

例题

树链剖分就是用两个dfs重新给树编号、深度、求父节点、求重儿子、子树大小、每条重链的顶端,然后用线段树维护新的编号的树

#include<iostream>
#include<algorithm>
using namespace std;
#define N 200010
int n,m,r,mod;
int head[N],nex[N],e[N],w[N],tot;
int dep[N],fa[N],siz[N],son[N],id[N],wt[N],top[N];
int cnt,res;
struct tree{
    int l,r,sum,lazy;
}t[N*4];


void add(int x,int y){
    e[++tot]=y,nex[tot]=head[x],head[x]=tot;
}

void push(int p){
    if(!t[p].lazy) return;
    t[p*2].lazy+=t[p].lazy;
    t[p*2+1].lazy+=t[p].lazy;
    t[p*2].sum+=t[p].lazy*(t[p*2].r-t[p*2].l+1);
    t[p*2+1].sum+=t[p].lazy*(t[p*2+1].r-t[p*2+1].l+1);
    t[p*2].sum%=mod;
    t[p*2+1].sum%=mod;
    t[p].lazy=0;
}

void build(int p,int l,int r){
    t[p].l=l,t[p].r=r;
    if(l==r){
        t[p].sum=wt[l];
        return ;
    }
    int mid = l+r>>1;
    build(2*p,l,mid);
    build(2*p+1,mid+1,r);
    t[p].sum=(t[2*p].sum+t[2*p+1].sum)%mod;
}

void update(int p,int l,int r,int k){
    if(l<=t[p].l&&r>=t[p].r){
        t[p].lazy+=k;
        t[p].sum+=(t[p].r-t[p].l+1)*k;
        return ;
    }
    int mid =t[p].l+t[p].r>>1;
    push(p);
    if(l<=mid) update(p*2,l,r,k);
    if(r>mid) update(p*2+1,l,r,k);
    t[p].sum=(t[2*p].sum+t[2*p+1].sum)%mod;
}

int query(int p,int l,int r){
    if(l<=t[p].l&&r>=t[p].r) return t[p].sum;
    int mid =t[p].l+t[p].r>>1;
    push(p);
    int sum=0;
    if(l<=mid) sum+=query(p*2,l,r);
    if(r>mid) sum+=query(p*2+1,l,r);
    return sum%mod;
}


//标记深度 子树大小 重儿子 父亲
void dfs1(int x,int f,int deep){
    dep[x]=deep;
    fa[x]=f;
    siz[x]=1;
    int maxson=-1;
    for(int i=head[x];i;i=nex[i]){
        int y=e[i];
        if(y==f) continue;
        dfs1(y,x,deep+1);
        siz[x]+=siz[y];
        if(siz[y]>maxson) son[x]=y,maxson=siz[y];
    }
}

//每个点的新编号 赋值每个点的初始值到新编号上 处理每条重链
void dfs2(int x,int topf){
    id[x]=++cnt;
    wt[cnt]=w[x];
    top[x]=topf;
    if(!son[x]) return ;
    dfs2(son[x],topf);
    for(int i=head[x];i;i=nex[i]){
        int y=e[i];
        if(y==fa[x]||y==son[x]) continue;
        dfs2(y,y);
    }
}//每一条重链的新编号是连续的
//x->y之间的路径
int qrange(int x,int y){
    int ans=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        ans+=query(1,id[top[x]],id[x]);//线段树求区间和
        ans%=mod;
        x=fa[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    ans+=query(1,id[x],id[y]);
    return ans%mod;
}

int qson(int x){
    return query(1,id[x],id[x]+siz[x]-1);
}

void uprange(int x,int y,int k){
    k%=mod;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        update(1,id[top[x]],id[x],k);
        x=fa[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    update(1,id[x],id[y],k);
}

void upson(int x,int k){
    update(1,id[x],id[x]+siz[x]-1,k);
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n>>m>>r>>mod;
    for(int i=1;i<=n;i++) cin>>w[i];
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        add(x,y);add(y,x);
    }
    dfs1(r,0,1);
    dfs2(r,r);
    build(1,1,n);
    for(int i=0;i<m;i++){
        int k,x,y,z;
        cin>>k;
        if(k==1){
            cin>>x>>y>>z;
            uprange(x,y,z);
        }
        else if(k==2){
            cin>>x>>y;
            cout<<qrange(x,y)<<endl;
        }
        else if(k==3){
            cin>>x>>y;
            upson(x,y);
        }
        else{
            cin>>x;
            cout<<qson(x)<<endl;
        }
    }
    return 0;
}
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值