Walking on the tree

Link: https://ac.nowcoder.com/acm/contest/11334/J
Source: Niuke

Title description
Niu Niu worked hard to practice martial arts skills-light water floating, and ultimately failed to practice, but he learned the ability to walk on trees.

On this day, Niuniu fell into the enemy’s trap. There was a boulder chasing behind him. There were n points in front of him, and n-1 edges connected to form a connected graph (a tree). Now Niuniu must immediately choose to enter this picture, but Niuniu found , There are two different points in this picture. Once you enter a point, all the points that are different from that point will disappear (the connected edges will disappear), Niu Niu can only go to the point where the edges are connected, Niu Niu wants to be himself Try to have as many points as possible, so which points can he enter?

Input description: The
first line has a positive integer n {}n which means there are n {}n points (n\leq2×10^5) (n≤2×10
5
) The
second line has n{}n numbers a_ia
i represent two types of point (0 \ Leq a_i \ Leq. 1) (0≤a i ≦ 1) the next n-1 {} n-1 lines each have two positive integers u, v (u, v\leq n)u,v(u,v≤n) means there is an edge between u{}u and v{}v. Output description: the number of points that can be entered in the first line of output

The second line outputs the numbers of these points from small to large

Example 1
Input
Copy
3
1 1 0
1 2
1 3
Output
Copy
2
1 2
Description
There can be 2 moving positions when falling into 1 and 2, which is the largest
Example 2
Input
Copy
4
1 1 0 0
1 2
2 3
3 4
Output
Copy
4
1 2 3 4
Description
No matter where it falls, there are 2 positions that can be moved

If the two points are of the same type, put them in the same set. Finally, see how many points are in the largest set, and finally see how many points are equal to the largest set.

#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;
}