牛客小白月赛28 J.树上行走

题目链接

题目描述

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

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

输入描述:

第一行有一个正整数 $n$ 表示共有 $n$ 个点 $（n\le 2\times 10^5）$
第二行有 $n$ 个数 $a_i$ 表示两种类型的点 $（0\le a_i\le 1）$
接下来 $n-1$ 行每行有两个正整数 $u，v（u，v\le n）$ 表示 $u$ 和 $v$ 之间有一条边

输出描述:

第一行输出可以进入的点的个数

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

示例1

输入

3
1 1 0
1 2
1 3

输出

2
1 2

示例2

输入

4
1 1 0 0
1 2
2 3
3 4

输出

4
1 2 3 4

典型的并查集~
题目首先要求最大连通块，我们可以在合并时顺带合并数量，合并操作也很简单，就是当两者种类相同时就合并，最后查找每个点所在联通块的大小，如果等于最大值就放到答案数组里，AC代码如下：

#include <bits/stdc++.h>
using namespace std;
const int N=2e5+5;
int n,u,v,mx=0,vis[N],type[N],father[N],sum[N];
vector<int>ans;
int Findfather(int x){
    
    
    return father[x]==x?x:father[x]=Findfather(father[x]);
}

void Union(int x,int y){
    
    
    x=Findfather(x),y=Findfather(y);
    if(x!=y) father[y]=x,sum[x]+=sum[y];
}

void init(){
    
    
    for(int i=0;i<=n;i++)
        father[i]=i,sum[i]=1;
}

int main() {
    
    
    cin>>n;
    init();
    for(int i=1;i<=n;i++) cin>>type[i];
    for(int i=1;i<n;i++){
    
    
        cin>>u>>v;
        if(type[u]==type[v]) Union(u,v);
        int x=Findfather(u),y=Findfather(v);
        mx=max({
    
    mx,sum[x],sum[y]});
    }
    for(int i=1;i<=n;i++){
    
    
        if(sum[Findfather(i)]==mx) ans.push_back(i);
    }
    cout<<ans.size()<<endl;
    for(auto i:ans) cout<<i<<" ";
    return 0;
}