#include "infection.h"
#include <vector>
#include <utility>
#include <queue>
#include <iostream>
#include  <cassert>
#define ll long long
#define pb push_back
const ll NMAX = 151000;
const ll INF = 1e9;
std::vector <ll> adj[NMAX + 1];
ll tata[NMAX + 1];

bool infected[NMAX + 1];
ll coada[2 * NMAX + 1], viz[NMAX], N, M;
void bfs(ll start, ll last){
    ll st =1, dr = 1;
    coada[st] = start;
    viz[start] = 1;
    while(st <= dr){
        for(ll i = 0;i<adj[coada[st]].size();i++){
            ll vec = adj[coada[st]][i];
            if(!viz[vec]){
                coada[++dr] = vec;
                viz[vec] = 1;
                tata[vec] = coada[st];
            }
        }
        st++;
    }
}
void reset(){
    for(ll i = 0;i<N;i++){
        viz[i] = 0;
    }
}
bool check(ll nod){
    if(infected[nod]){
       //  std::cout<<"In ziua: "<< ziua<<", am fost infectat, eu "<<crt<<" in orasul "<<nod<<std::endl;
        return true;
    }
    return false;
}
std::vector<int> solve(
    std::vector<std::pair<int, int>> R,
    std::vector<long long> D,
    std::vector<int> S,
    std::vector<int> T,
    std::vector<bool> I
) {
    N = R.size() + 1;
    M = D.size();
    std::vector<int> rez(M);
    for(ll i = 0;i<N - 1;i++){
        adj[R[i].first].pb(R[i].second);
        adj[R[i].second].pb(R[i].first);
    }

    for(ll i = 0;i<M;i++){
//        std::cout<<"Am inceput bfs ul de la"<<S[i]<<" la "<< T[i]<<std::endl;
        bfs(T[i], i);
        reset();
        if(I[i]){
        ll cur = S[i];
        ll ziua = 1;
        while(cur != T[i]){
            cur = tata[cur];
            ziua++;
        }
        rez[i] = ziua;
        }
    }
//    for(ll i = 0;i<M;i++){
//        std::cout<<"Pentru nodul "<<i<<". Printam drumul de la "<<S[i]<<" la "<<T[i]<<":\n";
//        ll cur = S[i];
//        while(cur != T[i]){
//            std::cout<<cur<<" ";
//            cur = tata[i][cur];
//        }
//    }
    for(ll i = 0;i<M;i++){
        assert(rez[i] >= 0);
    }
    return {rez};
}

