dp学习笔记

题目
题意: n个点的树，每个节点具有一种颜色。任选一个点作为根，使得所有以非根为节点的子树满足颜色一致。
思路: 找规律，如果存在某个点满足所有特殊边都连过他，即满足。特殊边是指连接不同颜色的边。
或者树形dp，但是难度有些高。先一遍dfs求出以1为根是否满足条件，之后再dfs。如果某个点的所有儿子均满足条件，可；否则，只有当有且仅有一个儿子不满足条件且当前节点与父节点同色，且其他儿子与当前节点同色且满足条件才行。日，打死俺也想不出来的换根dp.
时间复杂度: O(n)
代码:

// Problem: A. Timofey and a tree
// Contest: Codeforces - Codeforces Round #395 (Div. 1)
// URL: https://codeforces.com/problemset/problem/763/A
// Memory Limit: 256 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include<bits/stdc++.h>
using namespace std;
const int N = 1e5+10;
int n,m,k,T;
int a[N];
vector<int> va[N];
bool f[N]; //以i为根是否同色
void dfs(int cur,int fa)
{
    
	f[cur] = 1;
	for(int j:va[cur])
	{
    
		if(j==fa) continue;
		dfs(j,cur);
		if(!f[j]||a[cur]!=a[j]) f[cur]=0;
	}
}
int dfs2(int cur,int fa)
{
    
	// cout<<cur<<":"<<fa<<"\n";
	int cnt = 0; //有几个儿子满足条件
	int cnt2 = 0; //有几个儿子满足条件且与父节点同色，这样才能换根之后以父节点连的其他分支满足条件
	int wh = 0;
	for(int j:va[cur])
	{
    
		if(j==fa) continue;
		if(f[j]) 
		{
    
			cnt++;
			if(a[cur]==a[j]) cnt2++;
		}
		else wh = j;
	}
	int tot = va[cur].size();
	if(cur!=1) tot--;
	if(cnt == tot) return cur;
	if(cnt2 == tot-1 && a[cur]==a[fa])
	{
    
		int t = dfs2(wh,cur);
		if(t!=-1) return t;
	}
	return -1;
}
void solve()
{
    
	cin>>n;
	for(int i=0;i<n-1;++i)
	{
    
		int x,y; cin>>x>>y;
		va[x].push_back(y),va[y].push_back(x);
	}
	for(int i=1;i<=n;++i) cin>>a[i]; a[0] = a[1];
	dfs(1,0); 
	// for(int i=1;i<=n;++i) cout<<i<<":"<<f[i]<<"\n";
	int ans = dfs2(1,0);
	if(ans==-1) cout<<"NO";
	else cout<<"YES\n"<<ans;
}
signed main(void)
{
    
	T = 1;
	// cin>>T;
	while(T--)
	solve();
	return 0;
}