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

using namespace std;

typedef long long ll;

const int MAXN = 5005;
vector<int> con[MAXN];
int n, m;

int par[MAXN], depth[MAXN];
set<int> child[MAXN];
bool c[MAXN], nxt[MAXN];
int path[MAXN][MAXN]; // ind, curr_location
int dist[MAXN];
vector<int> loc_to[MAXN];

void dfs(int curr, int p){
    for (auto i : con[curr])
        if (i != p)
            par[i] = curr, depth[i] = depth[curr] + 1, dfs(i, curr);
    
    child[curr].insert(curr);
    for (auto i : con[curr])
        for (auto j : child[i])
            child[curr].insert(j);
}

vector<int> solve(vector<std::pair<int, int>> r, vector<ll> d, vector<int> s, vector<int> t, vector<bool> inf) {
    n = r.size() + 1, m = d.size();

    for (int i = 0; i < n - 1; i++)
        con[r[i].first].push_back(r[i].second), con[r[i].second].push_back(r[i].first);

    dfs(0, 0);

    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < n; j++)
            path[i][j] = -1;

        int curr = s[i];
        while (child[curr].find(t[i]) == child[curr].end())
            path[i][curr] = par[curr], curr = par[curr], dist[i]++;
        
        curr = t[i];
        int last;

        while (child[curr].find(s[i]) == child[curr].end())
            path[i][par[curr]] = curr, last = curr, curr = par[curr], dist[i]++;
        path[i][curr] = last;
    }

    ll time = 0;

    vector<pair<int, int>> times;
    for (int i = 0; i < m; i++)
        times.push_back({d[i], i});

    sort(times.begin(), times.end());

    queue<tuple<int, int, bool>> q;

    int ind = 0;
    vector<int> res(m, 0);

    auto push = [&](int curr, int loc, bool inf) {
        q.push({curr, loc, inf});
        if (inf)
            nxt[loc] = true;
    };

    while (ind < m or !q.empty()){

        int l = q.size();

        swap(nxt, c);
        for (int i = 0; i < n; i++)
            nxt[i] = false;

        for (int z = 0; z < l; z++){
            auto temp = q.front(); q.pop();

            int curr = get<0>(temp), loc = get<1>(temp);
            bool inf = get<2>(temp);

            if (c[loc])
                res[curr] = max(res[curr], dist[curr] - (int)(time - d[curr]) + 1);
            
            if (loc != t[curr]) {
                if (c[loc])
                    push(curr, path[curr][loc], true);
                else
                    push(curr, path[curr][loc], inf);
            }
        }

        time++;
        while (times[ind].first == time and ind < m) {
            int curr = times[ind].second;
            push(curr, s[curr], inf[curr]);
            ind++;
        }

        if (q.empty()){
            while (ind < m and (q.empty() or times[ind].first == time)) {
                int curr = times[ind].second;
                push(curr, s[curr], inf[curr]);
                time = times[ind].first;
                ind++;
            }
        }
    }

    return res;
}
