Lomsat gelral--CF600E (树上启发式合并)

题目描述：点击进入

题意

有一棵 n 个结点的以 1 号结点为根的有根树。
每个结点都有一个颜色，颜色是以编号表示的， i 号结点的颜色编号为 ci 。
如果一种颜色在以 x 为根的子树内出现次数最多，称其在以 x 为根的子树中占主导地位。显然，同一子树中可能有多种颜色占主导地位。
你的任务是对于每一个 i∈[1,n]，求出以 i 为根的子树中，占主导地位的颜色的编号和。

思路

在刘爷跟杨大佬的催促下我终于开始了这个“人均 dsu ”的学习
（这意味着我是人均之下，刘爷还诱惑我学串串， “ 居心叵测 ” 啊/(ㄒoㄒ)/~~）
一道练习树上启发式合并（dsu on tree）的经典题目
树启简介：
1、只有对子树的询问；
2、没有修改；
流程：
一、dfs 记录每个点的重儿子
二、dsu 遍历每一个节点;
1、 递归解决所有的轻儿子，同时消除递归产生的影响;
2、递归重儿子，不消除递归的影响;
3、统计所有轻儿子对答案的影响;
4、更新该节点的答案;
5、删除所有轻儿子对答案的影响;

代码

#include<iostream>
#include<cstdio>
#include<vector>
#include<cstring>
#include<set>
#define int long long
using namespace std;
const int maxn=1e5+10;
typedef long long ll;
typedef unsigned long long ull; 
int n,tot; 
int ans[maxn];//记录每个节点子树的答案 
int sum[maxn];// 子树的大小 
int col[maxn];//节点的颜色 
int cnt[maxn];//某种颜色的数量 
int fson;//标记重儿子 
int son[maxn];//统计节点的重儿子
int maxx;//记录子树中最多的颜色的数量
int p;//记录子树中最多的颜色的节点编号的和。 
int head[maxn]; 
struct node 
{
	int to;
	int next;
}edge[maxn<<1];
void add(int u,int v)
{
	edge[tot].to=v; 
	edge[tot].next=head[u]; 
	head[u]=tot++;
}
void dfs(int pos,int fa)//查找重儿子
{
	sum[pos]=1;
	for(int i=head[pos];i!=-1;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa) continue;
		dfs(v,pos);
		sum[pos]+=sum[v];
		if(sum[son[pos]]<sum[v]) son[pos]=v;//更新重儿子 
	}
}
void change(int pos,int fa,int d) 
{
	cnt[col[pos]]+=d;//当前节点的颜色数量 + d 
	if(cnt[col[pos]]>maxx)//更新最大颜色的数量以及编号和 
	{
		maxx=cnt[col[pos]];
		p=col[pos];
	}
	else if(cnt[col[pos]]==maxx) p+=col[pos];//更新编号和 
	for(int i=head[pos];i!=-1;i=edge[i].next)// 
	{
		int v=edge[i].to;
		if(v==fa||v==fson/*绕过被标记的重儿子*/) continue;
		change(v,pos,d);
	}
}
void dsu(int pos,int fa,int keep)//keep：是否保留当前子树的数据
{
	for(int i=head[pos];i!=-1;i=edge[i].next)//遍历轻儿子，删除递归影响 
	{
		int v=edge[i].to;
		if(v==fa||v==son[pos]) continue;
		dsu(v,pos,0);
	}
	if(son[pos])//遍历重儿子，保留影响 
	{
		dsu(son[pos],pos,1);
		fson=son[pos]; //标记重儿子
	}
	change(pos,fa,1); //把轻儿子与重儿子的数据合并
	fson=0;//取消标记
	ans[pos]=p; //更新记录答案 
	if(!keep)//删除所有轻儿子对答案的影响 
	{
		change(pos,fa,-1);
		p=0;
		maxx=0;
	} 
}
signed main()
{
	memset(head,-1,sizeof(head));
    scanf("%lld",&n);
    for(int i=1;i<=n;i++) scanf("%lld",&col[i]);
    for(int i=1;i<n;i++)
    {
  	    int x,y;
  	    scanf("%lld%lld",&x,&y);
  	    add(x,y);
		add(y,x);
    }
    dfs(1,0); //从根节点开始
	dsu(1,0,1);
    for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
    return 0;
}