ツリー上のヒューリスティックマージテンプレート(より良いブルートフォース)

テイクCF600E Lomsatのgelralを例として

#include <bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;
typedef long long ll;
int a[maxn],sz[maxn],son[maxn];
int n;
vector<int> G[maxn];
//遍历确定轻重儿子 
void dfs(int u,int fa){
	sz[u]=1;
	for(int i=0;i<G[u].size();++i){
		int v=G[u][i];
		if(v==fa) continue;
		dfs(v,u);
		sz[u]+=sz[v];
		if(sz[v]>sz[son[u]]) son[u]=v;
	}
}
ll cnt[maxn],sum,maxx,Son,ans[maxn];
void add(int u,int fa,int val){		//val=1加上贡献 val=-1消除贡献 
	cnt[a[u]]+=val;
	if(cnt[a[u]]>maxx) maxx=cnt[a[u]],sum=a[u];
	else if(cnt[a[u]]==maxx) sum+=a[u];
	for(int i=0;i<G[u].size();++i){
		int v=G[u][i];
		if(v==fa||v==Son) continue;    
		add(v,u,val);
	}
}
void dfs2(int u,int fa,int xp) {	//xp=1保存 xp=0 消除 
	for(int i=0;i<G[u].size();++i){
		int v=G[u][i];
		if(v==fa||v==son[u]) continue;	//先遍历轻儿子
		dfs2(v,u,0);
	}
	if(son[u]) dfs2(son[u],u,1),Son=son[u];	//遍历重儿子 
	add(u,fa,1); Son=0;
	ans[u]=sum;
	if(xp==0) {
		add(u,fa,-1);sum=0,maxx=0;
	}
}
signed main() {
	scanf("%d",&n);
	for(int i=1;i<=n;i++) scanf("%d",a+i);
	for(int i=1;i<n;i++) {
		int u,v; scanf("%d%d",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs(1,0);
	dfs2(1,0,1);
	for(int i=1;i<=n;i++) 
		printf("%lld ",ans[i]);
} 

 

おすすめ

転載: blog.csdn.net/qq_44132777/article/details/109595484