Codeforces600E Lomsat gelral

dsu on tree

题目链接

点我跳转

题目大意

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

解题思路

dsu on tree 模板题（详见注释）

AC_Code

#include<bits/stdc++.h>
#define int long long 
using namespace std;
const int N = 3e5 + 10; 
struct Edge{
	int nex , to;
}edge[N << 1];
int head[N] , TOT;
void add_edge(int u , int v) // 链式前向星建图
{
	edge[++ TOT].nex = head[u] ;
	edge[TOT].to = v;
	head[u] = TOT;
}
int sz[N];   // sz[u] 表示以 u 为根的子树大小 
int hson[N]; // hson[u] 表示 u 的重儿子 
int HH;      // HH 表示当前根节点的重儿子 
int c[N];    //表示每个节点的颜色 
int cnt[N];  // cnt[i]表示当前子树中，颜色 i 出现了多少次 
int ans[N];  // 每个点为根的答案 
int n , res , ma; 
// ma  表示出现次数最多的颜色出现的次数 
// res 表示出现次数最多的颜色的颜色值总和 （两颜色出现次数相同则都要算） 
void dfs(int u , int far)
{
	sz[u] = 1;
	for(int i = head[u] ; i ; i = edge[i].nex) // 链式前向星 
	{
		int v = edge[i].to;
		if(v == far) continue ;
		dfs(v , u); 
		sz[u] += sz[v];  
		if(sz[v] > sz[hson[u]]) hson[u] = v; // 选择 u 的重儿子 
	}
}
void calc(int u , int far , int val) // 统计答案 
{
	if(val == 1) 
	{
		cnt[c[u]] ++ ;
		if(cnt[c[u]] > ma) ma = cnt[c[u]] , res = c[u];
		else if(cnt[c[u]] == ma) res += c[u];
	} 
	else cnt[c[u]] -- ; 
    for(int i = head[u] ; i ; i = edge[i].nex)
    {
        int v = edge[i].to;
        if(v == far || v == HH) continue ; // 如果 v 是当前根节点的重儿子，则跳过
        calc(v , u , val);
    }
} 
void dsu(int u , int far , int op)  // op 等于0表示不保留信息，等于1表示保留信息 
{
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == hson[u]) continue ; // 如果 v 是重儿子或者父亲节点就跳过 
		dsu(v , u , 0);     // 先遍历轻儿子 ，op = 0 :轻儿子的答案不做保留 
	}
	if(hson[u]) dsu(hson[u] , u , 1) , HH = hson[u];
	// 轻儿子都遍历完了，如果存在重儿子，遍历重儿子（事实上除了叶子节点每个点都必然有重儿子）
	// op = 1 ， 保留重儿子的信息 
	// 当前是以 u 为根节点的子树，所以根节点的重儿子 HH = hson[u]
	calc(u , far , 1); // 再次遍历轻儿子统计答案
	ans[u] = res;      // 更新答案 
	HH = 0;			   // 遍历结束 ，即将返回父节点，所以取消标记 HH 
	if(!op) calc(u , far , -1) , ma = 0 , res = 0; // 如果 op = -1，则 u 对于它的父亲来说是轻儿子，不需要将信息传递给它的父亲 
}
signed main()
{
	cin >> n;
	for(int i = 1 ; i <= n ; i ++) cin >> c[i];
	for(int i = 2 ; i <= n ; i ++)
	{
		int u , v;
		cin >> u >> v;
		add_edge(u , v) , add_edge(v , u);
	}
	dfs(1 , 0);
	dsu(1 , 0 , 0);
	for(int i = 1 ; i <= n ; i ++) cout << ans[i] << " \n"[i == n];
	return 0;
}