题解 | 树上行走

用 BFS 加并查集的办法来维护每个可以连通的区块来找到节点最多的连通块即可

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
const ll N=2e5+5;
vector<vector<ll>>adj(N);
bool vis[N];
ll cmp=1;
struct UF{
    ll rank=0;
    UF* root;
    ll cnt=1;
    bool p;
};

UF a[N];

UF* Root(UF* a){
    if(a->root!=a){
        a->root=Root(a->root);
    }
    return a->root;
}

void merge(UF* rx,UF*ry){
    if(rx==ry)return;
    if(rx->rank<ry->rank){
        rx->root=ry;
        ry->cnt+=rx->cnt;
        cmp=max(ry->cnt,cmp);
    }else{
        ry->root=rx;
        rx->cnt+=ry->cnt;
        cmp=max(rx->cnt,cmp);
        if(rx->rank==ry->rank){
            rx->rank++;
        }
    }
}

void bfs(ll x){
    queue<ll>q;
    q.push(x);

    while(!q.empty()){
        ll t=q.front();
        q.pop();

        for(auto it:adj[t]){
            if(vis[it]==0&&a[t].p==a[it].p){
                merge(Root(&a[t]),Root(&a[it]));
                vis[it]=1;
                q.push(it);
            }
        }
    }
}


int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    ll n;
    cin>>n;
    for(ll i=1;i<=n;i++){
        cin>>a[i].p;
        a[i].root=&a[i];
    }
    for(ll i=1;i<=n-1;i++){
        ll u,v;
        cin>>u>>v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for(ll i=1;i<=n;i++){
        if(vis[i]==0){
            vis[i]=1;
            bfs(i);
        }
    }
    
    vector<ll>ans;
    for(ll i=1;i<=n;i++){
        if(cmp==Root(&a[i])->cnt){
            ans.push_back(i);
        }
    }
    cout<<ans.size()<<'\n';
    bool f=0;
    for(auto t:ans){
        if(f==0){
            cout<<t;
            f=1;
        }else cout<<' '<<t;
    }

    return 0;
}