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

using namespace std;

const int maxn = 5005;

int n, q;
vector<int> g[maxn];
int parent[maxn];
vector<int> path;

unordered_map<int, vector<int>> passingThrough[maxn];

void dfs(int node, int par)
{
    parent[node] = par;

    for (auto to : g[node])
    {
        if (to != par)
        {
            dfs(to, node);
        }
    }
}

struct Traveler
{
    long long day;
    int from, to;
    bool isSick;
    int idx;

    inline friend bool operator < (const Traveler& lf, const Traveler& rf)
    {
        if (lf.isSick != rf.isSick) return lf.isSick < rf.isSick;
        return lf.day < rf.day;
    }
};

Traveler a[maxn];
vector<int> answer;
int daysLeft[maxn][maxn];

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;

    for (auto [x, y] : R)
    {
        g[x].push_back(y);
        g[y].push_back(x);
    }

    q = D.size();
    answer.resize(q);

    for (int i = 0; i < q; ++i)
    {
        a[i] = {D[i], S[i], T[i], I[i], i};
        answer[i] = 0;
    }

    sort(a + 1, a + q + 1);

    for (int it = 0; it < q; ++it)
    {
        auto& [day, from, to, isSick, idx] = a[it];

        dfs(from, -1);

        vector<int> path;
        int node = to;

        while (node != -1)
        {
            path.push_back(node);
            node = parent[node];
        }

        reverse(path.begin(), path.end());
        /*cout << "Path" << endl;
        for (auto node : path) cout << node << " ";
        cout << endl;*/

        for (int i = 0; i < path.size(); ++i)
        {
            daysLeft[it][path[i]] = path.size() - i;
        }
    }

    for (int it = 0; it < q; ++it)
    {
        auto& [day, from, to, isSick, idx] = a[it];

        dfs(from, -1);

        vector<int> path;
        int node = to;

        while (node != -1)
        {
            path.push_back(node);
            node = parent[node];
        }

        reverse(path.begin(), path.end());
        /*cout << "Path" << endl;
        for (auto node : path) cout << node << " ";
        cout << endl;*/

        if (isSick == true)
        {
            answer[idx] = path.size();
            //cout << "set to " << path.size() << endl;
        }

        for (int i = 0; i < path.size(); ++i)
        {
            passingThrough[path[i]][day].push_back(it);

            if (!isSick)
            {
                for (auto other : passingThrough[path[i]][day])
                {
                    if (a[other].isSick == true)
                    {
                        isSick = true;
                        a[it].isSick = true;
                        answer[idx] = max(answer[idx], int(path.size()) - i);
                        break;
                    }
                }
            }

            if (isSick)
            {
                for (auto other : passingThrough[path[i]][day])
                {
                    if (a[other].isSick == false)
                    {
                        a[other].isSick = true;
                        answer[other] = max(answer[other], daysLeft[other][path[i]]);
                    }
                }
            }

            ++day;
        }
    }

    return answer;
}
