#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;
bool get_path(int u, int dest, int prev = -1) {
    path.push_back(u);
    if(u == dest) return true;
    for(int v : adj[u]) {
        if(v != prev) {
            bool x = get_path(v, dest, u);
            if(x) return true;
        }
    }
    path.pop_back();
    return false;
}

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>>> mp;
    for(int i = 0; i < m; i++) {
        int u = S[i];
        int v = T[i];
        ll time = D[i];
        infected[i] = I[i];
        path.clear();
        get_path(u, v);
        int d = int(path.size());
        dist[i] = d - 1;
        for(int j = 0; j < d; j++) {
            mp[time + j][path[j]].push_back(i);
        }
    }

    for(int i = 0; i < m; i++) {
        if(infected[i]) {
            infect_time[i] = D[i];
        }
    }

    for(auto& [time, mpp] : mp) {
        for(auto& [node, v] : mpp) {
            int cnt = 0;
            for(int x : v) {
                cnt += infected[x];
            }
            if(cnt > 0) {
                for(int x : v) {
                    if(!infected[x]) {
                        infected[x] = true;
                        infect_time[x] = time;
                    }
                }
            }
        }
    }

    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;
}
