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;
}