rt 悬关

题目

题目描述

给定一棵树，其n个顶点的编号为1到n。第i条边连接顶点a和顶点b。顶点i的颜色为c[i]。如果从顶点1到顶点i的路径上（最短路径），不包含与点i相同颜色的点，我们称顶点i为好点。请按照升序输出所有好点。

输入格式
第一行输入一个正整数 ，表示顶点的数量。第二行输入n个正整数 ，分别表示每个点的颜色。 之后  行，每行输入两个正整数 ，表示一条树边。

输出格式
使用换行符作为分隔符，按升序输出所有好点的编号。

我的代码 自测爆0了

#include<bits/stdc++.h>
using namespace std;
int n,c[100003],g[100003],vis[100003],cnt=2;
vector<int>son[100003];
void dfs(int a){
	if(son[a].size()==0) return; 
	for(int i=1;i<=son[a].size();i++){
		int co=son[a][i];
		if(!vis[c[co]]){
			vis[c[co]]=1;
			g[cnt]=co;
			cnt++;
			dfs(co);
			vis[c[co]]=0;
		}else return;
	}
}
int main(){
cin>>n;
for(int i=1;i<=n;i++) cin>>c[i];
for(int i=1;i<=n-1;i++){
int a,b;
cin>>a>>b;
son[a].push_back(b);
}
vis[1]=c[1];
g[1]=1;
dfs(1);
sort(g+1,g+n+1);
for(int i=1;i<=cnt;i++) cout<<g[i]<<"\n";
	return 0;
}