BF的資料結構題單-提高組——樹鏈剖分

行走天涯的豆沙包發表於2020-11-14

樹鏈剖分板子題

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define ls 2 * rt
#define rs 2 * rt + 1
#define N 100005
int n, m; vector<int> edge[N];
int id[N], nw[N], w[N], tot;
int dep[N], sz[N], top[N], fa[N], son[N],r,mod;
inline void add(int u, int v){ edge[u].push_back(v); edge[v].push_back(u); }
struct Segment_Tree{
    static const int maxn = 1e5 + 5;
    struct node{ int l, r; ll sum, add; }t[maxn << 2];
    inline void pushup(int rt){ t[rt].sum = t[ls].sum + t[rs].sum; }
    inline void pushdown(int rt){
        if(t[rt].add){
            t[ls].add = (t[ls].add+t[rt].add)%mod; t[ls].sum = (t[ls].sum+(t[ls].r - t[ls].l + 1) * t[rt].add)%mod;
            t[rs].add = (t[rs].add+t[rt].add)%mod; t[rs].sum = (t[rs].sum+(t[rs].r - t[rs].l + 1) * t[rt].add)%mod;
            t[rt].add = 0;
        }
    }
    inline void build(int rt, int l, int r){
        t[rt].l = l; t[rt].r = r; t[rt].add = 0;
        if(l == r){ t[rt].sum = nw[l]; return ; }
        int mid = (l + r) >> 1;
        build(ls, l, mid); build(rs, mid + 1, r);
        pushup(rt);
    }
    inline void update(int rt, int l, int r, int v){
        if(l <= t[rt].l && t[rt].r <= r){
            t[rt].add = (t[rt].add+v)%mod;
            t[rt].sum = (t[rt].sum+v * (t[rt].r - t[rt].l + 1))%mod;
            return ;
        }
        pushdown(rt);
        int mid = (t[rt].l + t[rt].r) >> 1;
        if(l <= mid) update(ls, l, r, v);
        if(r > mid) update(rs, l, r, v);
        pushup(rt);
    }
    inline ll query(int rt, int l, int r){
        if(l <= t[rt].l && t[rt].r <= r) return t[rt].sum;
        pushdown(rt);
        int mid = (t[rt].l + t[rt].r) >> 1; ll ans = 0;
        if(l <= mid) ans = (ans+query(ls, l, r))%mod;
        if(r > mid) ans = (ans+query(rs, l, r))%mod;
        return ans%mod;
    }
}seg;
void dfs1(int u, int f, int depth){
    dep[u] = depth; fa[u] = f; sz[u] = 1;
    for(auto &v : edge[u]){
        if(v == f) continue;
        dfs1(v, u, depth + 1);
        sz[u] += sz[v];
        if(sz[son[u]] < sz[v]) son[u] = v;
    }
}
void dfs2(int u, int f){
    id[u] = ++tot; nw[tot] = w[u]; top[u] = f;
    if(!son[u]) return ;
    dfs2(son[u], f);
    for(auto &v : edge[u]){
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}
inline void update_path(int u, int v, int k){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        seg.update(1, id[top[u]], id[u], k);
        u = fa[top[u]];
    }
    if(dep[u] < dep[v]) swap(u, v);
    seg.update(1, id[v], id[u], k);
}
inline ll query_path(int u, int v){
    ll ans = 0;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        ans = (ans+seg.query(1, id[top[u]], id[u]))%mod;
        u = fa[top[u]];
    }
    if(dep[u] < dep[v]) swap(u, v);
    ans = (ans+seg.query(1, id[v], id[u]))%mod;
    return ans;
}
inline void update_tree(int u, int k){
    seg.update(1, id[u], id[u] + sz[u] - 1, k);
}
inline ll query_tree(int u){
    return seg.query(1, id[u], id[u] + sz[u] - 1);
}
inline void solve(){
    scanf("%d%d%d%d", &n,&m,&r,&mod);
    for(int i = 1 ; i <= n ; i ++) scanf("%d", &w[i]);
    for(int i = 1, u, v ; i < n ; i ++) scanf("%d %d", &u, &v), add(u, v);
    dfs1(r, -1, 1);
    dfs2(r, r);
    seg.build(1, 1, n);
    while(m --){
        int opt, u, v, k; scanf("%d", &opt);
        if(opt == 1){
            scanf("%d %d %d", &u, &v, &k);
            update_path(u, v, k);
        }else if(opt == 3){
            scanf("%d%d",&u, &k);
            update_tree(u, k);
        }else if(opt == 2){
            scanf("%d%d",&u, &v);
            printf("%lld\n", query_path(u, v));
        }else if(opt == 4){
            scanf("%d",&u);
            printf("%lld\n", query_tree(u));
        }
    }
}
signed main(){
    solve();
#ifndef ONLINE_JUDGE
    cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
    return 0;
}

相關文章