洛谷 P3384 - 【模板】树链剖分 （点权）

https://www.luogu.org/problemnew/show/P3384

 

 

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <string.h>
using namespace std;
const int N = 1e5+66;
#define LL long long
#define lt x<<1
#define rt x<<1|1
int n,m,r,mod;
int head[N],nxt[N<<1],to[N<<1],cnt;//邻接表
int v[N];//old_点权
int tree[N<<2],lazy[N<<2];
int son[N],id[N],fa[N],dep[N],sz[N],top[N],val[N];
//son重儿子 id新编号 fa父亲节点 dep深度 sz子树大小 top重链的顶端 val新的点权
int tm;//dfs序


void add(int x,int y)
{
    to[cnt]=y;
    nxt[cnt]=head[x];
    head[x]=cnt++;
}

void pushdown(int x,int l,int r)
{
    lazy[lt]+=lazy[x];
    lazy[rt]+=lazy[x];
    int mid = (l+r)>>1;
    tree[lt]+=(mid-l+1)*lazy[x];
    tree[rt]+=(r-mid)*lazy[x];
    lazy[x]=0;
    lazy[lt]%=mod;lazy[rt]%=mod;
    tree[lt]%=mod;tree[rt]%=mod;
}

void up(int x,int l,int r)
{
    tree[x]=tree[lt]+tree[rt];
    tree[x]%=mod;
}

void build(int x,int l,int r)
{
    lazy[x]=0;
    tree[x]=val[l];
    if(l==r) return;
    int mid = (l+r)>>1;
    build(lt,l,mid);
    build(rt,mid+1,r);
    up(x,l,r);
}

int query(int x,int l,int r,int ll,int rr)
{
    int ans=0;
    if(ll<=l&&rr>=r){
        return tree[x];
    }
    if(lazy[x])
        pushdown(x,l,r);
    int mid = (l+r)>>1;
    if(ll<=mid) ans+=query(lt,l,mid,ll,rr);
    ans%=mod;
    if(mid<rr) ans+=query(rt,mid+1,r,ll,rr);
    ans%=mod;
    up(x,l,r);
    return ans;
}

void update(int x,int l,int r,int ll,int rr,int k)
{
    if(ll<=l&&rr>=r){
        lazy[x]+=k;
        tree[x]+=k*(r-l+1);
        tree[x]%=mod;lazy[x]%=mod;
        return;
    }
    if(lazy[x])
        pushdown(x,l,r);
    int mid =(l+r)>>1;
    if(ll<=mid) update(lt,l,mid,ll,rr,k);
    if(mid<rr) update(rt,mid+1,r,ll,rr,k);
    up(x,l,r);
}

//-------以上为线段树-----------

void dfs1(int x,int f,int deep)//x当前节点，f父亲，deep深度
{
    dep[x]=deep;//标记每个点的深度
    fa[x]=f;//标记每个点的父亲
    sz[x]=1;//标记每个非叶子节点的子树大小
    int maxson=-1;
    for(int i=head[x];~i;i=nxt[i]){
        int y=to[i];
        if(y==f) continue;
        dfs1(y,x,deep+1);
        sz[x]+=sz[y];
        if(sz[y]>maxson) son[x]=y,maxson=sz[y];//标记每个非叶子节点的重儿子编号
     }
}

void dfs2(int x,int topf)//x当前节点，topf当前链的最顶端的节点
{
    id[x]=++tm;//标记每个点的新编号
    val[tm]=v[x];//把每个点的初始值赋到新编号上来
    top[x]=topf;//这个点所在链的顶端
    if(!son[x]) return;//如果没有儿子则返回
    dfs2(son[x],topf);//按先处理重儿子，再处理轻儿子的顺序递归处理
    for(int i=head[x];~i;i=nxt[i]){
        int y=to[i];
        if(y==fa[x]||y==son[x]) continue;
        dfs2(y,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);//把x点改为所在链顶端的深度更深的那个点
        ans+=query(1,1,n,id[top[x]],id[x]);
        ans%=mod;
        x=fa[top[x]];//把x跳到x所在链顶端的那个点的上面一个点
    }
    if(dep[x]>dep[y]) swap(x,y);//把y点变为深度更深的那个点
    ans+=query(1,1,n,id[x],id[y]);
    ans%=mod;
    return ans;
}


void upRange(int x,int y,int k){//更新树上的两点最短路径上的点权，都加k
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        update(1,1,n,id[top[x]],id[x],k);
        x=fa[top[x]];
    }
    if(dep[x]>dep[y])swap(x,y);
    update(1,1,n,id[x],id[y],k);
}

int qson(int x)
{
    return query(1,1,n,id[x],id[x]+sz[x]-1);//子树区间右端点为id[x]+siz[x]-1
}

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


void init()
{
    tm=0;cnt=0;//dnf序和邻接表
    for(int i=1;i<=n;i++){
        head[i]=-1;
    }
    //线段树的初始化可以在build里
}

int main()
{
    scanf("%d%d%d%d",&n,&m,&r,&mod);
    init();
    for(int i=1;i<=n;i++) scanf("%d",&v[i]);

    for(int i=1;i<n;i++){
        int a,b;
        scanf("%d%d",&a,&b);
        add(a,b);add(b,a);
    }

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



    while(m--){
        int k,x,y,z;
        scanf("%d",&k);
        if(k==1){
            scanf("%d%d%d",&x,&y,&z);
            upRange(x,y,z);
        }
        // 1 x y z 表示将树从x到y结点最短路径上所有节点的值都加上z
        else if(k==2){
            scanf("%d%d",&x,&y);
            printf("%d\n",qRange(x,y));
        }
        // 2 x y 表示求树从x到y结点最短路径上所有节点的值之和
        else if(k==3){
            scanf("%d%d",&x,&y);
            upson(x,y);
        }
        // 3 x z 表示将以x为根节点的子树内所有节点值都加上z
        else{
            scanf("%d",&x);
            printf("%d\n",qson(x));
        }
        //4 x 表示求以x为根节点的子树内所有节点值之和
    }

    return 0;
}