#include "infection.h"
#include <vector>
#include <utility>
#include <map>
#include <stdio.h>
using ll = long long;

const int MAX_N = 5000;
const int MAX_M = 5000;
std::vector<int> adj[MAX_N];
bool infected[MAX_M];
ll infect_time[MAX_M];
int dist[MAX_M];

std::vector<int> path;
int dfs(int u, int t, int prev = 0, int dist = 0) {
    path.push_back(u);
    if(u == t) {
        return dist;
    }
    for(int v : adj[u]) {
        if(v != prev) {
            int x = dfs(v, t, u, dist + 1);
            if(x > -1) {
                return x;
            }
        }
    }
    path.pop_back();
    return -1;
}

std::vector<int> solve(
    std::vector<std::pair<int, int>> R,
    std::vector<ll> D,
    std::vector<int> S,
    std::vector<int> T,
    std::vector<bool> I
) {
    int n = int(R.size()) + 1;
    for(int i = 0; i < n; i++) {
        adj[i].clear();
    }
    for(int i = 0; i < n - 1; i++) {
        auto [u, v] = R[i];
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    int m = int(D.size());
    std::map<ll, std::map<int, std::vector<int>>> events;
    for(int id = 0; id < m; id++) {
        int u = S[id];
        int v = T[id];
        ll time = D[id];
        path.clear();
        int distance = dfs(u, v);
        dist[id] = distance;
        for(int i = 0; i < distance + 1; i++) {
            events[time + i][path[i]].push_back(id);
        }
    }

    for(int i = 0; i < m; i++) {
        infected[i] = I[i];
        if(infected[i]) {
            infect_time[i] = D[i];
        } else {
            infect_time[i] = -1;
        }
    }
    for(auto& [t, mp] : events) {
        for(auto& [node, ids] : mp) {
            int cnt = 0;
            for(int id : ids) {
                cnt += infected[id];
            }
            if(cnt > 0) {
                for(int id : ids) {
                    if(!infected[id]) {
                        infect_time[id] = t;
                        infected[id] = true;
                    }
                }
            }
        }
    }

    std::vector<int> sol;
    for(int i = 0; i < m; i++) {
        if(!infected[i]) {
            sol.push_back(0);
        } else {
            sol.push_back((D[i] + dist[i]) - infect_time[i] + 1);
        }
    }

    return sol;
}
