CF600E Lomsat gelral

CF600E Lomsat gelral

• (dsu on tree) 入门例题.讲解.
• 注意 (dfs) 中各个操作的顺序,这样能保证清空轻儿子贡献时不会对重儿子造成影响.

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define pii pair<int,int>
inline int read()
{
	int x=0;
	bool pos=1;
	char ch=getchar();
	for(;!isdigit(ch);ch=getchar())
		if(ch=='-')
			pos=0;
	for(;isdigit(ch);ch=getchar())
		x=x*10+ch-'0';
	return pos?x:-x;
}
const int MAXN=1e5+10;
vector<int> sons[MAXN];
int n;
int mxson[MAXN],siz[MAXN],col[MAXN],cnt[MAXN];
ll sum,ans[MAXN];
void getsiz(int u,int fa)
{
	siz[u]=1;
	int tot=sons[u].size();
	for(int i=0;i<tot;++i)
		{
			int v=sons[u][i];
			if(v!=fa)
				{
					getsiz(v,u);
					if(siz[v]>siz[mxson[u]])
						mxson[u]=v;
					siz[u]+=siz[v];
				}
		}
}
int mxcnt,bigson;
void add(int u,int fa,int val)
{
	cnt[col[u]]+=val;
	if(cnt[col[u]]>mxcnt)
		mxcnt=cnt[col[u]],sum=col[u];
	else if(cnt[col[u]]==mxcnt)
		sum+=col[u];
	int tot=sons[u].size();
	for(int i=0;i<tot;++i)
		{
			int v=sons[u][i];
			if(v!=fa && v!=bigson)
				add(v,u,val);
		}
}
void dfs(int u,int fa,bool keep)
{
	int tot=sons[u].size();
	for(int i=0;i<tot;++i)
		{
			int v=sons[u][i];
			if(v!=fa && v!=mxson[u])
				dfs(v,u,false);
		}
	if(mxson[u])
		dfs(mxson[u],u,true),bigson=mxson[u];
	add(u,fa,1);
	ans[u]=sum;
	bigson=0;
	if(keep==false)
		add(u,fa,-1),sum=0,mxcnt=0;
}
int main()
{
	n=read();
	for(int i=1;i<=n;++i)	
		col[i]=read();
	for(int i=1;i<n;++i)
		{
			int u=read(),v=read();
			sons[u].push_back(v);
			sons[v].push_back(u);		
		}
	getsiz(1,0);
	dfs(1,0,true);
	for(int i=1;i<=n;++i)
		printf("%I64d ",ans[i]);
	return 0;
}