[算法学习] dsu on tree
简介
dsu on tree跟dsu没有关系,但是dsu on tree借鉴了dsu的启发式合并的思想。
它是用来解决一类树上的询问问题,一般这种问题有以下特征:
\(1.\)只有对子树的查询;
\(2.\)没有修改。
如果满足以上特征,那么dsu on tree很可能就可以派上用场了。
算法
我们以CF600E Lomsat gelral为例。
Descrption
一棵以\(1\)为根的树有 \(n\) 个结点,每个结点的颜色是\(c_i\),每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。
数据范围\(1\le n\le 10^5, 1\le c_i\le 10^5\)。
Solution
我们考虑暴力怎么写:遍历每一个节点 \(u\) ,然后把子树内的所有颜色暴力统计出来更新答案。
然后消除节点 \(u\) 的贡献,继续递归计算其他点的贡献。
复杂度\(O(n^2)\),显然是很不优美的。
然后dsu on tree就登场了。它也是一个暴力,但是它结合了轻重链剖分,将复杂度降到了\(O(nlogn)\)。
我们先跑一个\(dfs1\),记录节点\(u\)的重儿子\(heavy[u]\)。
接下来:
- 遍历每一个节点。
- 递归所有的轻儿子,递归进入时暴力加贡献,递归结束时暴力消除贡献。
- 递归重儿子,递归进入时暴力加贡献,递归结束时不消除贡献。
- 统计所有轻儿子的答案,并更新该节点的答案。
- 删除所有轻儿子对答案的影响。
复杂度\(O(nlogn)\),可以通过本题。
Code
// Author: wlzhouzhuan
#pragma GCC optimize(2)
#pragma GCC optimize(3)
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define ull unsigned long long
#define rint register int
#define rep(i, l, r) for (rint i = l; i <= r; i++)
#define per(i, l, r) for (rint i = l; i >= r; i--)
#define mset(s, _) memset(s, _, sizeof(s))
#define pb push_back
#define pii pair <int, int>
#define mp(a, b) make_pair(a, b)
inline int read() {
int x = 0, neg = 1; char op = getchar();
while (!isdigit(op)) { if (op == '-') neg = -1; op = getchar(); }
while (isdigit(op)) { x = 10 * x + op - '0'; op = getchar(); }
return neg * x;
}
inline void print(int x) {
if (x < 0) { putchar('-'); x = -x; }
if (x >= 10) print(x / 10);
putchar(x % 10 + '0');
}
const int N = 100005;
int n;
vector <int> adj[N];
void add(int u, int v) { adj[u].pb(v); }
int sz[N], fa[N], heavy[N]; // heavy[u] 表示u的重儿子
void dfs1(int u, int f) {
sz[u] = 1, fa[u] = f;
for (auto v: adj[u]) {
if (v == f) continue;
dfs1(v, u);
sz[u] += sz[v];
if (sz[v] > sz[heavy[u]]) {
heavy[u] = v; // 更新重儿子
}
}
}
long long ans[N], sum;
int cnt[N], col[N], Son, Max;
void change(int u, int val) {
cnt[col[u]] += val;
if (cnt[col[u]] > Max) Max = cnt[col[u]], sum = col[u];
else if (cnt[col[u]] == Max) sum += col[u];
for (auto v: adj[u]) {
// 由于重儿子的信息没有被删去,所以已经统计过了,不能再计算
if (v == fa[u] || v == Son) continue;
change(v, val);
}
}
void dfs2(int u, int keep) {
for (auto v: adj[u]) {
if (v == fa[u] || v == heavy[u]) continue;
dfs2(v, 0); // 遍历u的轻儿子
}
if (heavy[u]) dfs2(heavy[u], 1), Son = heavy[u];
change(u, 1), ans[u] = sum, Son = 0;
if (!keep) change(u, -1), sum = Max = 0;
}
int main() {
n = read();
for (rint i = 1; i <= n; i++) {
col[i] = read();
}
for (rint i = 1; i < n; i++) {
int x = read(), y = read();
add(x, y), add(y, x);
}
int root = 1; // 题目默认1为根
dfs1(root, 0);
dfs2(root, 0);
for (rint i = 1; i <= n; i++) {
printf("%lld ", ans[i]);
}
return 0;
}