洛谷P3384 【模板】轻重链剖分

洛谷P3384 【模板】轻重链剖分

传送门[https://www.luogu.com.cn/problem/P3384]

方法

树链划分+线段树

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int maxn = 1e5+1000;
int f[maxn], head[maxn], mod, rt, len, dfn[maxn], n, m, son[maxn], size[maxn], a[maxn], top[maxn], dep[maxn], cnt, rk[maxn];
struct edge{ int to, next; } e[2*maxn];
struct node{int l, r, w, siz, lazy; } tr[maxn<<2];
void add(int x, int y){ e[++len].to=y; e[len].next=head[x]; head[x]=len; }
void dfs1(int u, int fa){ 
    size[u] = 1; 
    for(int i=head[u]; ~i; i=e[i].next){
        int v = e[i].to;
        if(v == fa) continue;
        dep[v] = dep[u] + 1, f[v] = u;
        dfs1(v, u);
        size[u] += size[v];
        if(!son[u] || size[v] > size[son[u]]) son[u] = v;
    }
}
void dfs2(int u, int tp){
    top[u] = tp;
    dfn[u] = ++cnt;
    rk[cnt] = a[u];
    if(son[u]) dfs2(son[u], tp);
    for(int i=head[u]; ~i; i=e[i].next){
        int v = e[i].to;
        if(v!=son[u] && v!=f[u]) dfs2(v, v);
    }
}
void pushup(int u){ tr[u].w = (tr[u<<1].w + tr[u<<1|1].w + mod) % mod; }
void build(int u, int l, int r){
    tr[u].l = l; tr[u].r = r; tr[u].siz = r - l + 1;
    if(l==r){ tr[u].w = rk[l]; return; }
    int mid = (l+r)>>1;
    build(u<<1, l, mid); build(u<<1|1, mid+1, r);
    pushup(u);
}
void pushdown(int u){
    if(tr[u].lazy){
        tr[u<<1].w = (tr[u<<1].w + tr[u<<1].siz * tr[u].lazy)%mod;
        tr[u<<1|1].w = (tr[u<<1|1].w + tr[u<<1|1].siz * tr[u].lazy)%mod;
        tr[u<<1].lazy = (tr[u<<1].lazy + tr[u].lazy)%mod;
        tr[u<<1|1].lazy = (tr[u<<1|1].lazy + tr[u].lazy)%mod;
        tr[u].lazy=0;
    }
}
void update(int u, int l, int r, int w){
    if(l<=tr[u].l && r>=tr[u].r){
        tr[u].w += tr[u].siz * w;
        tr[u].lazy += w;
        return;
    } 
    pushdown(u);
    int mid = (tr[u].l + tr[u].r) >> 1;
    if(l<=mid) update(u<<1, l, r, w);
    if(r>mid)  update(u<<1|1, l, r, w);
    pushup(u);
}
void treeadd(int u, int v, int w){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        update(1, dfn[top[u]], dfn[u], w);
        u = f[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    update(1, dfn[u], dfn[v], w);
}
int query(int u, int l, int r){
    int ans = 0;
    if(l<=tr[u].l && r>=tr[u].r) return tr[u].w;
    pushdown(u);
    int mid = (tr[u].l + tr[u].r)>>1;
    if(l<=mid) ans = (ans + query(u<<1, l, r))%mod;
    if(r>mid)  ans = (ans + query(u<<1|1, l, r))%mod;
    return ans;
}
void querysum(int u, int v){
    int ans = 0;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        ans = (ans + query(1, dfn[top[u]], dfn[u]))%mod;
        u = f[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    ans = (ans + query(1, dfn[u], dfn[v]))%mod;
    printf("%d\n", ans);
}
int main(){
    memset(head, -1, sizeof(head));
    scanf("%d%d%d%d", &n, &m, &rt, &mod);
    for(int i=1; i<=n; i++) scanf("%d", &a[i]);
    for(int i=1; i<n; i++){
        int x, y; scanf("%d%d", &x, &y);
        add(x, y); add(y, x);
    }
    dfs1(rt, 0);
    dfs2(rt, rt);
    build(1, 1, n);
    while(m--){
        int op, x, y, z; scanf("%d", &op);
        if(op==1){
            scanf("%d%d%d", &x, &y, &z);
            treeadd(x, y, z%mod);
        }else if(op==2){
            scanf("%d%d", &x, &y);
            querysum(x, y);
        }else if(op==3){
            scanf("%d%d", &x, &z);
            update(1, dfn[x], dfn[x]+size[x]-1, z%mod);
        }else if(op==4){
            scanf("%d", &x);
            printf("%d\n", query(1, dfn[x], dfn[x]+size[x]-1));
        }
    }
    return 0;
}