Codeforces #395 Div1 A. Timofey and a tree

いろいろグダグダでレート減。
http://codeforces.com/contest/763/problem/A

問題

N頂点の木を成すグラフが与えられる。
各頂点には色が塗られている。

ここから以下の条件を満たす頂点の有無を判定し、存在するなら1例を示せ。

• その頂点を根とした根付き木を構築したとき、各子頂点のsubtreeは同一色の頂点だけからなる。

解法

頂点の色で考えるのではなく、辺を考えるとよい。
異なる色同士の頂点を結ぶ辺があったとする。
問題の条件を満たすには、そのような辺が根頂点に接続されていてはならない。

よって、辺のうち異なる色の頂点を結ぶ辺を数え上げ、各頂点についてそれらすべてにつながる頂点があればそれが解。

int N;

int U[101010],V[101010];
int C[101010];
vector<int> E[101010];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>U[i]>>V[i];
		U[i]--;
		V[i]--;
		E[U[i]].push_back(V[i]);
		E[V[i]].push_back(U[i]);
	}
	FOR(i,N) cin>>C[i];
	
	int num=0;
	FOR(i,N-1) if(C[U[i]]!=C[V[i]]) num++;
	FOR(i,N) {
		int x=0;
		FORR(e,E[i]) if(C[i]!=C[e]) x++;
		if(x==num) {
			cout<<"YES"<<endl;
			cout<<i+1<<endl;
			return;
		}
	}
	cout<<"NO"<<endl;
	
}

まとめ

Bの方がラクじゃないですかね…。