CodeForces-600E Lomsat gelral DSU on Tree 模板題

題意

• 有一顆\(n\)個結點，以1為根的有根樹
• 每個結點有一種顏色，顏色以編號表示，\(i\)號結點的顏色編號為\(c_i\)
• 如果一種顏色以\(x\)為根的子樹內出現最多，稱其為\(x\)為根的子樹中占主導地位。顯然同一個子樹中可能有多種顏色占主導地位。
• 求出\(1-n\)的子樹中，占主導地位的顏色的編號和

分析

典型的啟發式合併題(DSU on Tree)

核心思想：利用重鏈剖分的性質優化子樹貢獻的計算

來解決一類不帶修改的子樹查詢問題

• 預處理重兒子
• dfs所有輕兒子，並清空貢獻
• dfs重兒子，不清空貢獻
• 暴力合併除了重兒子以外的貢獻

程式碼

int fa[maxn],son[maxn],dep[maxn],siz[maxn],top[maxn];
vector<int> e[maxn];
ll a[maxn],cnt[maxn],ans[maxn];
ll sum;
int mx,Son;

void dfs1(int u){
	siz[u] = 1;
	for(auto v:e[u]){
		if(v == fa[u]) continue;
		dep[v] = dep[u] + 1;
		fa[v] = u;
		dfs1(v);
		siz[u] += siz[v];
		if(siz[v] > siz[son[u]]) 
			son[u] = v;
	}
}

void add(int x,int val){
	cnt[a[x]] += val;
	if(cnt[a[x]] > mx) mx = cnt[a[x]],sum = a[x];
	else if(cnt[a[x]] == mx) sum += a[x];
	for(auto v:e[x]){
		if(v == fa[x] || v == Son) continue;
		add(v,val);
	}
}

void dfs(int x,int op){
	for(auto v:e[x]){
		if(v == fa[x]) continue;
		if(v != son[x]) dfs(v,0);
	}
	if(son[x]) dfs(son[x],1),Son = son[x];
	add(x,1);
	Son = 0;
	ans[x] = sum;
	if(!op) add(x,-1),sum = 0,mx = 0;
}

int main(){
	int n = readint();
	for(int i = 1;i <= n;i++)
		a[i] = readint();
	for(int i = 1;i < n;i++){
		int u = readint();
		int v = readint();
		e[u].pb(v);
		e[v].pb(u);
	}
	dfs1(1);
	dfs(1,0);
	for(int i = 1;i <= n;i++) 
		cout << ans[i] << ' ';
}