#include "infection.h"
#include <vector>
#include <utility>
#include <map>
#include <stdio.h>
#include <algorithm>
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];

struct Event {
    ll time;
    int node;
    int id;
};

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::vector<Event> events;
    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++) {
            events.push_back({time + j, path[j], i});
        }
    }

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

    std::sort(events.begin(), events.end(), [](const Event& a, const Event& b) {
        if(a.time != b.time) {
            return a.time < b.time;
        }
        if(a.node != b.node) {
            return a.node < b.node;
        }
        return a.id < b.id;
    });

    int i = 0;
    int sz = int(events.size());
    while(i < sz) {
        int j = i;
        ll time = events[i].time;
        int node = events[i].node;
        std::vector<int> v;
        while(j < sz && events[j].time == time && events[j].node == node) {
            v.push_back(events[j].id);
            j++;
        }
        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;
                }
            }
        }
        i = j;
    }

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