题目链接:https://codeforces.com/problemset/problem/600/E
一直没有点这个技能点,今天跟队友打训练赛,碰到一道 \(dsu\ on\ tree\) 的题写不出来,就回来把这个题写了
\(dsu\ on\ tree\) 运用了轻重链剖分的思想,先处理轻儿子的答案,然后消去轻儿子的影响,最后处理重儿子,并保留重儿子的答案,
时间复杂度 \(O(nlogn)\)
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<cmath>
using namespace std;
typedef long long ll;
const int maxn = 200010;
int n;
int h[maxn], cnt = 0;
struct E{
int to, next;
}e[maxn << 1];
void add(int u, int v){
e[++cnt].next = h[u];
e[cnt].to = v;
h[u] = cnt;
}
ll sum, ans[maxn];
int mx, Son, sz[maxn], son[maxn], c[maxn], cn[maxn * 10];
void dfs1(int u, int par){
sz[u] = 1;
int mx = 0;
for(int i = h[u] ; i != -1; i = e[i].next){
int v = e[i].to;
if(v == par) continue;
dfs1(v, u);
if(sz[v] > mx){
mx = sz[v];
son[u] = v;
}
sz[u] += sz[v];
}
}
void update(int u, int par, int val){
cn[c[u]] += val;
if(cn[c[u]] > mx) mx = cn[c[u]], sum = c[u];
else if(cn[c[u]] == mx) sum += c[u];
for(int i = h[u] ; i != -1 ; i = e[i].next){
int v = e[i].to;
if(v == par || v == Son) continue;
update(v, u, val);
}
}
void dfs2(int u, int par, int is){ // is: 是不是重儿子
// 处理当前子树的轻儿子的答案,处理完后轻儿子的答案被消除
for(int i = h[u] ; i != -1 ; i = e[i].next){
int v = e[i].to;
if(v == par || v == son[u]) continue;
dfs2(v, u, 0);
}
// 再处理重儿子 重儿子的答案保留
sum = 0, mx = 0;
if(son[u]) dfs2(son[u], u, 1);
Son = son[u];
update(u, par, 1); Son = 0; // 计算当前子树的答案,跳过保留答案的重儿子,只暴力统计轻儿子的答案
// for(int i = 1 ; i <= n ; ++i){
// printf("%d ", cn[i]);
// } printf("\n");
ans[u] = sum;
if(!is) update(u, par, -1); // 如果是父亲的轻儿子,就消去当前子树的影响
}
ll read(){ ll s = 0, f = 1; char ch = getchar(); while(ch < '0' || ch > '9'){ if(ch == '-') f = -1; ch = getchar(); } while(ch >= '0' && ch <= '9'){ s = s * 10 + ch - '0'; ch = getchar(); } return s * f; }
int main(){
memset(h, -1, sizeof(h));
n = read();
for(int i = 1 ; i <= n ; ++i){
c[i] = read();
}
int u, v;
for(int i = 1 ; i < n ; ++i){
u = read(), v = read();
add(u, v), add(v, u);
}
dfs1(1, 0);
// for(int i = 1 ; i <= n ; ++i) printf("%d ", son[i]); printf("\n");
dfs2(1, 0, 1);
for(int i = 1 ; i <= n ; ++i){
printf("%lld ", ans[i]);
} printf("\n");
return 0;
}