树链剖分（模板）洛谷P3384

喜欢ヅ旅行 2021-10-30 05:12 287阅读 0赞

#include<iostream>
    #include<cstdio>
    #include<cstring>
    #include<algorithm>
    
    #define lson rt<<1, l, mid
    #define rson rt<<1|1, mid+1, r
    using namespace std;
    const int maxn = 500100;
    int P;
    typedef long long ll;
    
    struct Edge {
        int to, next;
    }edges[maxn*2];
    int head[maxn], tot;
    int n;
    ll w[maxn];
    int fa[maxn], son[maxn], siz[maxn], dep[maxn];
    int top[maxn], id[maxn], rk[maxn], cnt;
    // son[u]: u的重儿子
    // top[u]: u所在链的顶端节点
    // 
    // id[u]: 树链剖分后节点的新编号
    // rk[u]: dfs编号对应的节点  rk[id[u]] = u
    // cnt : dfs序
    
    void add(int u, int v) {
        edges[++tot].to = v;
        edges[tot].next = head[u];
        head[u] = tot;
    }
    
    // 求fa, siz, dep, son
    void dfs1(int u) { 
        siz[u] = 1;
        for(int i=head[u];i;i=edges[i].next) {
            int v = edges[i].to;
            if(v==fa[u]) continue;
    
            fa[v] = u;
            dep[v] = dep[u] + 1;
            dfs1(v);
            siz[u] += siz[v];
    
            if(siz[v]>siz[son[u]])
                son[u] = v;
        }
    }
    
    // 连接重链
    void dfs2(int u, int topf) {
        id[u] = ++cnt;
        rk[cnt] = u;
    
        top[u] = topf;
        if(!son[u])
            return;
    
        dfs2(son[u], topf);    // 重链先dfs，保证重链上各节点dfs序连续
    
        for(int i=head[u];i;i=edges[i].next) {
            int v = edges[i].to;
            if(v!=fa[u] && v!=son[u])  // 不是重儿子，轻链
                dfs2(v, v);
        }
    }
    
    /*
    操作1： 格式： 1 x y z 表示将树从x到y结点最短路径上所有节点的值都加上z
    
    操作2： 格式： 2 x y 表示求树从x到y结点最短路径上所有节点的值之和
    
    操作3： 格式： 3 x z 表示将以x为根节点的子树内所有节点值都加上z
    
    操作4： 格式： 4 x 表示求以x为根节点的子树内所有节点值之和
     */
    
    
    ll sum[maxn*4], lazy[maxn*4];
    
    void build(int rt, int l, int r) {
        if(l!=r) {
            int mid = (l+r)/2;
            build(lson);
            build(rson);
            sum[rt] = (sum[rt<<1] + sum[rt<<1|1]) % P;
        } else {
            sum[rt] = w[rk[l]];
            lazy[rt] = 0;
        }
    }
    void pushDown(int rt, int len) {
        if(lazy[rt]) {
            (sum[rt<<1] += (len-(len>>1)) * lazy[rt] % P) %= P; 
            (sum[rt<<1|1] += (len>>1) * lazy[rt] % P) %= P;
    
            (lazy[rt<<1] += lazy[rt]) %= P;
            (lazy[rt<<1|1] += lazy[rt]) %= P;
            lazy[rt] = 0;
        }
    }
    
    void update(int rt, int l, int r, int L, int R, ll add) {
        if(L<=l && R>=r) {
            (lazy[rt] += add) %= P;
            (sum[rt] += (r-l+1)*add) %= P;
            return;
        }
    
        int mid = (l+r)/2;
        pushDown(rt, r-l+1);
        
        if(L<=mid)
            update(lson, L, R, add);
        if(R>mid)
            update(rson, L, R, add);
    
        sum[rt] = (sum[rt<<1] + sum[rt<<1|1]) % P;
    }
    
    ll query(int rt, int l, int r, int L, int R) {
        if(L<=l && R>=r) return sum[rt];
    
        ll res = 0;
        int mid = (l+r)/2;
        pushDown(rt, r-l+1);
    
        if(L<=mid) res += query(lson, L, R) % P;
        if(R>mid) res += query(rson, L, R) % P;
        return res % P;
    }
    
    // 操作1
    void updatePath(int x, int y, ll z) {
        while(top[x]!=top[y]) {
            if(dep[top[x]]<dep[top[y]]) swap(x, y);
                update(1, 1, n, id[top[x]], id[x], z);
            x = fa[top[x]];
        }
        if(dep[x]>dep[y]) swap(x, y);
        update(1, 1, n, id[x], id[y], z);
    }
    
    
    // 操作3
    void updateSon(int x, ll z) {
        update(1, 1, n, id[x], id[x]+siz[x]-1, z);
    }
    
    
    // 操作2
    ll queryPath(int x, int y) {
        ll res = 0;
        while(top[x]!=top[y]) {
            if(dep[top[x]]<dep[top[y]]) swap(x, y);
                res = (res + query(1, 1, n, id[top[x]], id[x])) % P;
            x = fa[top[x]];
        }
        if(dep[x]>dep[y]) swap(x, y);
        res = (res + query(1, 1, n, id[x], id[y])) % P;
        return res;
    }
    
    // 操作4
    ll qeurySon(int x) {
        return query(1, 1, n, id[x], id[x]+siz[x]-1);
    }
    
    int main() {
        int M, R;
        cin>>n>>M>>R>>P;
        for(int i=1;i<=n;i++) {
            scanf("%lld", &w[i]);
        }
        for(int i=0;i<n-1;i++) {
            int u, v;
            scanf("%d %d", &u, &v);
            add(u, v);
            add(v, u);
        }
        
        dfs1(R);
        dfs2(R, R);
    
        build(1, 1, n);
    
        while(M--) {
            int op, x, y;
            ll z;
            scanf("%d", &op);
            if(op==1) {
                scanf("%d %d %lld", &x, &y, &z);
                updatePath(x, y, z);
            } else if(op==2) {
                scanf("%d %d", &x, &y);
                printf("%lld\n", queryPath(x, y));
            } else if(op==3) {
                scanf("%d %lld", &x, &z);
                updateSon(x, z);
            } else {
                scanf("%d", &x);
                printf("%lld\n", qeurySon(x));
            }
    
        }
        return 0;
    }

转载于:https://www.cnblogs.com/izcat/p/11375653.html