树上启发式合并 ( DSU on Tree) 是一种优雅的暴力
时间复杂度是 (O(nlogn))
启发式就是基于直觉或经验的意思
树上启发式合并的代码很简单
void dfs(int u,int p,bool keep){
for(u.lightson v in u.son){
dfs(v,u,false);
}
if(have u.heavyson)dfs(u.heavyson,u,true);;
count(); //统计 u 及所有轻儿子的贡献,并计算 u 的答案
if(keep == false) del(); //如果不保留的话,删除记录
}
这段代码多看几遍,自己多想一想
对于当前的节点
u
,在执行count()
计算节点u
的答案的时候,其所有子节点的答案都已经被计算出来,并且只有u
的重儿子及其子树的记录没有被删除,所以计算u
的答案的时候只要暴力遍历所有的轻儿子及其子树就可以。极其暴力(bushi
CF600E
首先要明白怎么统计答案
int mp[N],mx,sum; //mp是颜色出现次数,mx是最多出现的次数,sum是当前的最优节点的和
mp[color]++;
if (mp[color] > mx){
mx = mp[color];
sum = color;
}
else if (mp[color] == mx) sum += color;
然后就是开始暴力合并(启发式)
ll ans[N], sum, mx, mp[N];
void count(int u, int p,int hson) { //暴力算轻儿子及子树的答案,此时重儿子及其子树的信息还在。
mp[color[u]]++;
if (mp[color[u]] > mx) {
mx = mp[color[u]];
sum = color[u];
}
else if (mp[color[u]] == mx) sum += color[u];
repE(i, u) {
int v = E[i].to;
if (v == p or v == hson)continue;
count(v, u, hson);
}
}
void del(int u, int p) { //删除子树的信息。
mp[color[u]]--;
repE(i, u) {
int v = E[i].to;
if (v == p)continue;
del(v, u);
}
}
void dfs(int u, int p, bool keep) {
repE(i, u) {
int v = E[i].to;
if (v == p or v == son[u])continue;
dfs(v, u, false); //
}
if (son[u])dfs(son[u], u, true);
count(u, p, son[u]); ans[u] = sum;
if (not keep) {
sum = mx = 0;
del(u,p);
}
}
完整代码
#include<bits/stdc++.h>
#define repE(i,u) for(int i = head[u]; ~i;i = E[i].next)
using namespace std;
const int N = 1e5 + 10;
typedef long long ll;
int fa[N], sz[N], dep[N], son[N], color[N];
struct Edge { int to, next; }E[N << 1];
int head[N], tot;
void addEdge(int from, int to) { E[tot] = { to,head[from] }; head[from] = tot++; }
void dfs1(int u, int p) {
fa[u] = p;
dep[u] = dep[p] + 1;
sz[u] = 1;
int mx = -1;
repE(i, u) {
int v = E[i].to;
if (v == p)continue;
dfs1(v, u);
sz[u] += sz[v];
if (sz[v] > mx)mx = sz[v], son[u] = v;
}
}
ll ans[N], sum, mx, mp[N];
void count(int u, int p,int hson) {
mp[color[u]]++;
if (mp[color[u]] > mx) {
mx = mp[color[u]];
sum = color[u];
}
else if (mp[color[u]] == mx) sum += color[u];
repE(i, u) {
int v = E[i].to;
if (v == p or v == hson)continue;
count(v, u, hson);
}
}
void del(int u, int p) {
mp[color[u]]--;
repE(i, u) {
int v = E[i].to;
if (v == p)continue;
del(v, u);
}
}
void dfs(int u, int p, bool keep) {
repE(i, u) {
int v = E[i].to;
if (v == p or v == son[u])continue;
dfs(v, u, false);
}
if (son[u])dfs(son[u], u, true);
count(u, p, son[u]); ans[u] = sum;
if (not keep) {
sum = mx = 0;
del(u,p);
}
}
int n;
int main() {
ios::sync_with_stdio(0);
cin >> n; memset(head, -1, sizeof head);
for (int i = 1; i <= n; i++)cin >> color[i];
for (int i = 1; i < n; i++) {
int a, b; cin >> a >> b;
addEdge(a, b); addEdge(b, a);
}
dfs1(1, 0);
dfs(1, 0, true);
for (int i = 1; i <= n; i++)cout << ans[i] << "
"[i == n];
}