树上行走

链接：https://ac.nowcoder.com/acm/contest/11334/J
来源：牛客网

题目描述
牛牛苦练武功绝学——轻功水上漂，最终没有练成，但是他学会了在树上行走的本领。

这天，牛牛落入了敌人的陷阱，身后有巨石追击，面前有n个点，n-1条边连成一张连通图（一棵树），现在牛牛必须立马选择进入这张图中，但是牛牛发现，这张图有两种不同的点，一旦进入一个点，所有与该点不同类型的点都会消失（相连的边也会消失），牛牛只能走到有边相连的点，牛牛想要自己尽量有更多的点可以活动，那么他可以进入哪些点？

输入描述:
第一行有一个正整数 n {}n 表示共有 n {}n 个点（n\leq2×10^5）（n≤2×10
5
）
第二行有n{}n 个数 a_ia
i
​
表示两种类型的点（0 \leq a_i\leq 1）（0≤a
i
​
≤1）
接下来 n-1{}n−1行每行有两个正整数u，v（u，v\leq n）u，v（u，v≤n） 表示 u{}u 和 v{}v 之间有一条边
输出描述:
第一行输出可以进入的点的个数

第二行从小到大输出这些点的编号

示例1
输入
复制
3
1 1 0
1 2
1 3
输出
复制
2
1 2
说明
落到1和2的情况可以有2的移动位置，是最大的
示例2
输入
复制
4
1 1 0 0
1 2
2 3
3 4
输出
复制
4
1 2 3 4
说明
不论落到哪个点，都有2个位置可以移动

如果两个点的类型是一样的话，就把他们放到同一个集合里面，最后看最大的集合个数是多少，最后看有多少点的集合和最大集合的个数是相等的

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>

using namespace std;

const int N = 2e5 + 10;

int p[N], a[N], cnt[N];

int find(int x){
	if (x != p[x])   p[x] = find(p[x]);
	return p[x];
}

int main(){
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	
	int n;
	scanf("%d", &n);
	
	for (int i = 1; i <= n; i ++){
		  p[i] = i;
		  cnt[i] = 1;
	}
	
	for (int i = 1; i <= n; i ++){
		scanf("%d", &a[i]);
	}
	
	for (int i = 1; i < n; i ++){
		int x, y;
		scanf("%d%d", &x, &y);
		if (a[x] != a[y])    continue;
		else{
			int pa = find(x);
			int pb = find(y);
		//	cout << "---------" << endl;
			if (pa != pb){
				   p[pa] = pb;
				   cnt[pb] += cnt[pa];
			}
		}
	}
	
	int ans = 0;
	int cnt1 = 0;
	for (int i  = 1; i <= n; i ++){
		int p = find(i);
	//	cout << "----" << endl;
		ans = max(cnt[p], ans);
	}
	
	vector<int> v;
	for (int i = 1; i <= n; i ++){
		int p = find(i);
		if (cnt[p] == ans){
			v.push_back(i);
		}
	}
	
	cout << v.size() << endl;
	
	for (int i = 0; i < v.size(); i ++){
		cout << v[i] << " ";
	}
	
	cout << endl;
	
	return 0;
}

标签：cnt,include,cout,int,牛牛,行走,树上,find
来源： https://blog.csdn.net/qq_45772483/article/details/112558164