#include "infection.h"
#include <bits/stdc++.h>

using namespace std;
using ll = long long;
const int NMAX = 5001;

vector<int> g[NMAX];
vector<int> drum;

bool det_drumu(int nod, int f, int dr){
    drum.push_back(nod);
    if(nod == dr){
        return 1;
    }
    for(const int &x : g[nod]){
        if(x == f) continue;
        if( det_drumu(x, nod, dr) ) return 1;
    }
    drum.pop_back();
    return 0;
}

vector<int> solve(
    vector<pair<int, int>> r,
    vector<long long> d,
    vector<int> s,
    vector<int> t,
    vector<bool> infect
) {

    int n = r.size() + 1;
    for(int i = 0; i < n - 1; i++){
        g[ r[i].first ].push_back(r[i].second);
        g[ r[i].second ].push_back(r[i].first);
    }

    int m = d.size();

    map< pair<ll, int>, vector<int> > mp; // cine e la momentul ala

    vector< pair<ll, int> > dr[m];
    for(int i = 0; i < m; i++){
        drum.clear();
        det_drumu(s[i], -1, t[i]);
        for(int j = 0; j < drum.size(); j++){
            dr[i].push_back({d[i] + j, drum[j]});
            mp[ dr[i].back() ].push_back(i);
        }
    }

    priority_queue< pair<int, int>, vector< pair<int, int> >, greater< pair<int, int> > > pq; // {cand e infectat, cine e infectat}
    bool done[n];
    vector<int> sol(m);
    for(int i = 0; i < n; i++) done[i] = sol[i] = 0;

    for(int i = 0; i < m; i++){
        if(infect[i]){
            pq.push({ d[i], i });
        }
    }

    while(!pq.empty()){
        int x = pq.top().second;
        ll tmp = pq.top().first; pq.pop();
        if(done[x]) continue; // l-am infectat deja
        sol[x] = dr[x].back().first - tmp + 1;
        done[x] = 1;

        for(const pair<ll, int> &y : dr[x]){
            if(y.first < tmp) continue;
            // infecteaza pe toti cei de la y
            for(const int &el : mp[y]){
                if(!done[el]){
                    pq.push({y.first, el});
                }
            }
        }
    }

    return sol;
}