369C - Valera and Elections - CodeForces Solution

#include <bits/stdc++.h>
#include <string.h>
#define ll long long
#define Fast_Warak ios_base::sync_with_stdio(false); cin.tie(NULL);
#define loop(n) for(ll i=0 ; i<n ; i++)
#define rloop(n) for(ll i=n-1 ; i>=0 ; i--)
#define lp(start, end) for(ll i=start ; i<=end ; i++)
#define rlp(start, end) for(ll i=start ; i>=end ; i--)
#define b_e(v) v.begin(), v.end()
 
const ll max_ele = 1e18;
 
using namespace std;
 
void file(){
  freopen("jumping.in", "r", stdin);
  freopen("output.txt", "w", stdout);
}
 
vector<vector<ll>>v;
vector<ll>visited;
vector<ll>parent;
 
void sizee(ll n){
  v.resize(n+1);
  visited.resize(n+1, -1);
  parent.resize(n+1);
}
 
void BFS(ll node){
  parent[node]=-1;
  visited[node]=0;
  queue<ll>q;
  q.push(node);
  while(!q.empty()){
    ll point=q.front();
    q.pop();
    for(int child:v[point]){
       if(visited[child]==-1){
        visited[child]=visited[point]+1;
        parent[child]=point;
        q.push(child);
       }
    }
  }
}
 
int main()
{
    // file()
    Fast_Warak
 
    int t_case=1;
    //cin >> t_case;
    
    while(t_case--)
    {
       ll n; cin >> n;
       sizee(n); ll arr[n+1]={0};
       priority_queue<pair<ll, ll>>q;
       loop(n-1){
        ll a, b, t; cin >> a >> b >> t;
        t--; if(t)arr[a]=t, arr[b]=t;
        v[a].push_back(b);
        v[b].push_back(a);
       }
       BFS(1);
       lp(1, n){
        if(arr[i]){
          q.push({visited[i], i});
        }
       }
       set<ll> s; vector<bool>vis(n+1);
       while(!q.empty()){
        pair<ll, ll> p=q.top();
        q.pop(); bool flag=false;
        while(p.second!=-1){
           if(!vis[p.second]&&arr[p.second]&&!flag){ 
              s.insert(p.second);
              flag=true;
           }
           if(vis[p.second]){
            s.erase(p.second);
            break;
           }
           vis[p.second]=true;
           if(p.second==1)break;
           p={vis[parent[p.second]], parent[p.second]};
        }
       }
       cout << s.size() << "\n";
       for(auto p:s)cout << p << " ";
    } 
    
    return 0;
}