送你个树链剖分板子(洛谷P3384 【模板】轻重链剖分)

送你个树链剖分板子(洛谷P3384 【模板】轻重链剖分)

题目链接

码量还是差了点,昨天一个错误找了一晚上都没找到,这块终于算入了个门吧

#include
#define lowbit(x) x&(-x)
#define ll long long
#define inf 0x3f3f3f3f
using namespace std;
const int maxn=1e5+5;
struct Edge{
    int to,nxt;
}edge[maxn<<2];   //无向图所以得乘4
int head[maxn],tot;
struct node{
    int l,r,v,len,lazy;
}tree[maxn<<2];
int n,m,root,mod;  //见题意
int dep[maxn];    //当前结点的深度
int siz[maxn];    //当前结点为根的子树结点个数和
int id[maxn];     //dfs序表示的新编号
int rk[maxn];     //dfs序编号对应在树的结点
int fa[maxn];    //当前结点的父亲
int son[maxn];   //保存重儿子
int top[maxn];   //重链的顶点
int cnt;
int a[maxn];    //初始结点的值

/**以下---向前星操作*/
void init(){
    memset(head,-1,sizeof(head));
    tot=0;
}
void add(int u,int v){
    edge[++tot].to=v;
    edge[tot].nxt=head[u];
    head[u]=tot;
}
/**以上---向前星操作*/

/**以下---线段树操作*/
void build(int now,int l,int r){
    tree[now].l=l,tree[now].r=r,tree[now].len=r-l+1,tree[now].lazy=0;
    if(l==r){
        tree[now].v=rk[l];
        return ;
    }
    int mid=(l+r)>>1;
    build(now<<1,l,mid);
    build((now<<1)+1,mid+1,r);
    tree[now].v=(tree[now<<1].v+tree[(now<<1)+1].v)%mod;
}
void push_down(int now){
    if(tree[now].lazy){
        tree[now<<1].lazy+=tree[now].lazy;
        tree[now<<1].lazy%=mod;
        tree[(now<<1)+1].lazy+=tree[now].lazy;
        tree[(now<<1)+1].lazy%=mod;
        tree[now<<1].v+=(tree[now<<1].len*tree[now].lazy);
        tree[(now<<1)].v%=mod;
        tree[(now<<1)+1].v+=(tree[(now<<1)+1].len*tree[now].lazy);
        tree[(now<<1)+1].v%=mod;
        tree[now].lazy=0;
    }
}
void update(int now,int l,int r,int val){
    if(tree[now].l>=l&&tree[now].r<=r){
        tree[now].v+=val*tree[now].len;
        tree[now].v%=mod;
        tree[now].lazy+=val;
        tree[now].lazy%=mod;
        return ;
    }
    if(tree[now].l>r||tree[now].r<l)
        return ;
    if(tree[now].lazy){
        push_down(now);
    }
    update(now<<1,l,r,val);
    update((now<<1)+1,l,r,val);
    tree[now].v=(tree[now<<1].v+tree[(now<<1)+1].v)%mod;
}
int query(int now,int l,int r){
    if(tree[now].l>=l&&tree[now].r<=r){
        return tree[now].v;
    }
    if(tree[now].l>r||tree[now].r<l)
        return 0;
    if(tree[now].lazy){
        push_down(now);
    }
    return (query(now<<1,l,r)+query((now<<1)+1,l,r))%mod;
}
/**以上---线段树操作*/

/**以下---两次dfs*/
void dfs1(int now,int father,int depth){
    //标记每个结点的父亲,深度,子树的结点数,以及每个结点的重儿子
    dep[now]=depth,fa[now]=father,siz[now]=1;
    int maxnum=-1;
    for(int i=head[now];i!=-1;i=edge[i].nxt){
        int v=edge[i].to;
        if(v==father)
            continue;
        dfs1(v,now,depth+1);
        siz[now]+=siz[v];
        if(siz[v]>maxnum){
            maxnum=siz[v];
            son[now]=v;
        }
    }
}
void dfs2(int now,int nowt){
    //连接重链,并且根据重链标注出dfs序,方便使用数据结构维护
    top[now]=nowt;
    id[now]=++cnt;
    rk[cnt]=a[now];
    if(son[now]==0){
        return ;
    }
    dfs2(son[now],nowt);
    for(int i=head[now];i!=-1;i=edge[i].nxt){
        int v=edge[i].to;
        if(v!=son[now]&&v!=fa[now])
            dfs2(v,v);
    }
}
/**以上---两次dfs*/

void updatePath(int x,int y,int val){
   while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])
            swap(x,y);
        update(1,id[top[x]],id[x],val);
        x=fa[top[x]];
   }
   if(dep[x]>dep[y])
        swap(x,y);
   update(1,id[x],id[y],val);
}
int queryPathSum(int x,int y){
    int ans=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])
            swap(x,y);
        ans=(ans+query(1,id[top[x]],id[x]))%mod;
        x=fa[top[x]];
    }
    if(dep[x]>dep[y])
        swap(x,y);
    ans=(ans+query(1,id[x],id[y]))%mod;
    return ans;
}
int main(){
    scanf("%d%d%d%d",&n,&m,&root,&mod);
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
    }
    init();   //向前星初始化
    int u,v;
    for(int i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    cnt=0;   //新编号
    dfs1(root,0,1);
    dfs2(root,root);
    build(1,1,n);

    while(m--){
        int flag,x,y,z;
        scanf("%d",&flag);
        if(flag==1){
            scanf("%d%d%d",&x,&y,&z);
            updatePath(x,y,z);
        }
        else if(flag==2){
            scanf("%d%d",&x,&y);
            printf("%d\n",queryPathSum(x,y));
        }
        else if(flag==3){
            scanf("%d%d",&x,&y);
            y%=mod;
            update(1,id[x],siz[x]+id[x]-1,y);
        }
        else{
            scanf("%d",&x);
            printf("%d\n",query(1,id[x],id[x]+siz[x]-1));
        }
    }
    return 0;
}

你可能感兴趣的:(数据结构与算法)