洛谷 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;
}

结果