洛谷 P3178 [HAOI2015] 树上操作-提高+/省选-

P3178 [HAOI2015] 树上操作

题目描述

有一棵点数为 $N$ 的树，以点 $1$ 为根，且树有点权。然后有 $M$ 个操作，分为三种：

• 操作 $1$：把某个节点 $x$ 的点权增加 $a$。
• 操作 $2$：把某个节点 $x$ 为根的子树中所有点的点权都增加 $a$。
• 操作 $3$：询问某个节点 $x$ 到根的路径中所有点的点权和。

输入格式

第一行包含两个整数 $N,M$。表示点数和操作数。
接下来一行 $N$ 个整数，表示树中节点的初始权值。
接下来 $N-1$ 行每行两个正整数 $from,to$， 表示该树中存在一条边 $(from,to)$。
再接下来 $M$ 行，每行分别表示一次操作。其中第一个数表示该操作的种类，之后接这个操作的参数。

输出格式

对于每个询问操作，输出该询问的答案。答案之间用换行隔开。

输入输出样例 #1

输入 #1

5 5
1 2 3 4 5
1 2
1 4
2 3
2 5
3 3
1 2 1
3 5
2 1 2
3 3

输出 #1

6
9
13

说明/提示

对于 $100\%$ 的数据，$1\le N,M\le 10^5$，且所有输入数据的绝对值都不会超过 $10^6$。

solution

• 1 通过树链剖分将操作中树的节点的 dfn 映射成连续的若干段
• 2 通过线段树处理区间修改和查询。
  第1步可以参考洛谷 P3384 【模板】重链剖分/树链剖分-提高+/省选-

代码

#include <iostream>
#include "bit"
#include "vector"
#include "unordered_set"
#include "set"
#include "queue"
#include "algorithm"
#include "bitset"
#include "cstring"

using namespace std;

/*
 * P3178 [HAOI2015] 树上操作
 * 题目大意：
 * 有一棵点数为 N 的树，以点 1 为根，且树有点权。然后有 M 个操作，分为三种，1≤N,M≤10^5：
 * 操作 1：把某个节点 x 的点权增加 a。
 * 操作 2：把某个节点 x 为根的子树中所有点的点权都增加 a。
 * 操作 3：询问某个节点 x 到根的路径中所有点的点权和。
 *
 * 思路：
 * 1 通过树链剖分将操作中树的节点的 dfn 映射成连续的若干段
 * 2 通过线段树处理区间修改和查询。
 */

typedef pair<int, int> pii;
typedef long long ll;

const int N = 1e5 + 5;

int n, m, val[N], d[N], dfn[N], top[N], id[N], t, fa[N], siz[N], son[N];
vector<int> e[N];

void dfs(int u, int p) {
    fa[u] = p;
    d[u] = d[p] + 1;
    siz[u] = 1;
    int Max = -1;
    for (int v: e[u]) {
        if (v == p) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if (siz[v] > Max) Max = siz[v], son[u] = v;
    }
}

void dfs2(int u, int tp) {
    dfn[u] = ++t;
    id[t] = u;
    top[u] = tp;
    if (!son[u]) return;
    dfs2(son[u], tp);

    for (int v: e[u]) {
        if (v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}

ll tag[N << 2], sum[N << 2];

void push_up(int rt) {
    sum[rt] = sum[rt << 1] + sum[rt << 1 | 1];
}

void push_down(int rt, int l, int r) {
    int mid = l + r >> 1;
    sum[rt << 1] += tag[rt] * (mid + 1 - l);
    sum[rt << 1 | 1] += tag[rt] * (r - mid);
    tag[rt << 1] += tag[rt];
    tag[rt << 1 | 1] += tag[rt];
    tag[rt] = 0;
}

void build(int l, int r, int rt) {
    if (l == r) {
        sum[rt] = val[id[l]];
        return;
    }

    int mid = l + r >> 1;
    build(l, mid, rt << 1);
    build(mid + 1, r, rt << 1 | 1);
    push_up(rt);
}

void update(int L, int R, int l, int r, int rt, int v) {
    if (L <= l && r <= R) {
        sum[rt] += 1ll * (r - l + 1) * v;
        tag[rt] += v;
        return;
    }

    push_down(rt, l, r);
    int mid = l + r >> 1;
    if (L <= mid) update(L, R, l, mid, rt << 1, v);
    if (R > mid) update(L, R, mid + 1, r, rt << 1 | 1, v);
    push_up(rt);
}

ll search(int L, int R, int l, int r, int rt) {
    if (L <= l && r <= R) return sum[rt];
    push_down(rt, l, r);
    int mid = l + r >> 1;
    ll s = 0;
    if (L <= mid) s += search(L, R, l, mid, rt << 1);
    if (R > mid) s += search(L, R, mid + 1, r, rt << 1 | 1);
    return s;
}

int main() {
    cin >> n >> m;

    for (int i = 1; i <= n; i++) cin >> val[i];
    for (int i = 1, x, y; i < n; i++) cin >> x >> y, e[y].push_back(x), e[x].push_back(y);

    dfs(1, 0);
    dfs2(1, 1);

    build(1, n, 1);
    for (int i = 1, op, x, a; i <= m; i++) {
        cin >> op >> x;
        if (op == 1) {
            cin >> a;
            update(dfn[x], dfn[x], 1, n, 1, a);
        } else if (op == 2) {
            cin >> a;
            update(dfn[x], dfn[x] + siz[x] - 1, 1, n, 1, a);
        } else {
            ll s = 0;
            while (x) {
                s += search(dfn[top[x]], dfn[x], 1, n, 1);
                x = fa[top[x]];
            }
            cout << s << endl;
        }
    }
    return 0;
}

结果