树上操作

一本通·提高篇 · 代码 · cpp · 难度 4/5 · 共 3213 字
以1为根的树上，支持单点加、子树加、查询某点
一本通
练习
正文

/*
原题：T1561 「一本通 4.5 练习 1」树上操作
题意：以 1 为根的树上，支持单点加、子树加、查询某点到根路径上所有点的点权和。
思路：HLD 把子树映射到 dfn 连续区间 [dfn[x], dfn[x]+sz[x]-1]，单点即 dfn[x]。
线段树维护区间和，支持区间加(lazy)与区间求和；点到根路径用 HLD 分段查询。
复杂度：时间 O((N+M)log^2 N)，空间 O(N)。
易错点：根固定为 1；子树区间长度是 sz[x] 而非整段；点到根路径要 HLD 分段跳转，
不能简单用前缀 [1,dfn[x]]（兄弟节点会被误算）。
*/
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;
int n, m;
vector<int> g[MAXN];
int a[MAXN];
int fa[MAXN], dep[MAXN], sz[MAXN], son[MAXN], top[MAXN], dfn[MAXN], rk[MAXN], tim;
void dfs1(int u, int p){
    fa[u] = p;
    dep[u] = dep[p] + 1;
    sz[u] = 1;
    son[u] = 0;
    for(int v : g[u]){
        if(v == p) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if(sz[v] > sz[son[u]]) son[u] = v;
    }
}
void dfs2(int u, int tp){
    top[u] = tp;
    dfn[u] = ++tim;
    rk[tim] = u;
    if(son[u]) dfs2(son[u], tp);
    for(int v : g[u]){
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}
long long sumv[4 * MAXN];
long long lazy[4 * MAXN];
void pushup(int o){
    sumv[o] = sumv[o * 2] + sumv[o * 2 + 1];
}
void pushdown(int o, int l, int r){
    if(lazy[o]){
        int mid = (l + r) / 2;
        sumv[o * 2] += lazy[o] * (mid - l + 1);
        lazy[o * 2] += lazy[o];
        sumv[o * 2 + 1] += lazy[o] * (r - mid);
        lazy[o * 2 + 1] += lazy[o];
        lazy[o] = 0;
    }
}
void build(int o, int l, int r){
    if(l == r){
        sumv[o] = a[rk[l]];
        return;
    }
    int mid = (l + r) / 2;
    build(o * 2, l, mid);
    build(o * 2 + 1, mid + 1, r);
    pushup(o);
}
void update(int o, int l, int r, int ql, int qr, long long v){
    if(ql <= l && r <= qr){
        sumv[o] += v * (r - l + 1);
        lazy[o] += v;
        return;
    }
    pushdown(o, l, r);
    int mid = (l + r) / 2;
    if(ql <= mid) update(o * 2, l, mid, ql, qr, v);
    if(qr > mid) update(o * 2 + 1, mid + 1, r, ql, qr, v);
    pushup(o);
}
long long qres;
void query(int o, int l, int r, int ql, int qr){
    if(ql <= l && r <= qr){
        qres += sumv[o];
        return;
    }
    pushdown(o, l, r);
    int mid = (l + r) / 2;
    if(ql <= mid) query(o * 2, l, mid, ql, qr);
    if(qr > mid) query(o * 2 + 1, mid + 1, r, ql, qr);
}
long long path_sum(int u){
    long long res = 0;
    while(top[u] != 1){
        qres = 0;
        query(1, 1, n, dfn[top[u]], dfn[u]);
        res += qres;
        u = fa[top[u]];
    }
    qres = 0;
    query(1, 1, n, dfn[1], dfn[u]);
    res += qres;
    return res;
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> m;
    for(int i = 1; i <= n; i++) cin >> a[i];
    for(int i = 1; i < n; i++){
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 1);
    build(1, 1, n);
    while(m--){
        int op;
        cin >> op;
        if(op == 1){
            int x, val;
            cin >> x >> val;
            update(1, 1, n, dfn[x], dfn[x], val);
        } else if(op == 2){
            int x, val;
            cin >> x >> val;
            update(1, 1, n, dfn[x], dfn[x] + sz[x] - 1, val);
        }else{
            int x;
            cin >> x;
            cout << path_sum(x) << "\n";
        }
    }
    return 0;
}
树上操作

正文

一本通·提高篇的其它内容
