A. Timofey and a tree_思維

題目大意：

給一顆無根樹，現要選出一個節點作為根，使得其所有子樹的節點顏色分別相同。問是否可能。

思路和程式碼：

一開始看到只有單組資料，試了試暴力做法，即將每個點都試做根，跑一遍搜尋。但是這樣做會T。

bool dfs(int now , int pre , vct<vct<int>> &eg , vct<int> &c){
	int f = 1 ;
	for(int nxt : eg[now]){
		if(nxt == pre) continue ;
		if(c[now] != c[nxt]) return 0 ;
		if(!dfs(nxt , now , eg , c)) return 0 ;
	}
	return 1 ; 
}

void solve(){
	int n ;
	cin >> n ;
	vct<int> c(n + 1 , 0) ;
	vct<vct<int> > eg(n + 1) ;
	rep(i , 2 , n){
		int u , v ; cin >> u >> v ;
		eg[u].pb(v) ;
		eg[v].pb(u) ;
	} 
	
	rep(i , 1 , n) cin >> c[i] ;
	 
	rep(i , 1 , n){
		int f = 1 ;
		for(auto rt : eg[i]){
			if(!dfs(rt , i , eg , c))f = 0 ;
			if(!f) break ;
		}
		
		if(f){
			cout << "YES\n" << i << "\n" ;
			return ;
		}
	} 
	cout << "NO\n" ;
	
}//code_by_tyrii 
 

於是想其他做法，發現要滿足所有子樹顏色分別相同，就只能把顏色不同的邊全部插在大根上。

所以程式碼也是比較簡單的，如下。

void solve2(){
	int n ;
	cin >> n ;
	
	vct<int> c(n + 1 , 0) ;
	vct<pii> e(n , {0 , 0}) ;
	vct<int> cnt(n + 1 , 0) ;
	int all = 0 ;
	
	rep(i , 0 , n - 2)
	cin >> e[i].fi >> e[i].se ;
	
	rep(i , 1 , n) cin >> c[i] ;
	
	rep(i , 0 , n - 1){
		int u = e[i].fi ; 
		int v = e[i].se ; 
		if(c[u] != c[v]){
			all ++ ;
			cnt[u] ++ , cnt[v] ++ ;
		}
	}
	rep(i , 1 , n){
		if(cnt[i] == all){
			cout << "YES\n" << i << "\n" ;
			return ;
		}
	}cout << "NO\n" ;
	
}