#include "infection.h"
#include <vector>
#include <utility>
#include <iostream>
#include <map>
#include <set>
using namespace std;

vector<int> path;
vector<long long> time_end;
set<long long> all_times;
vector<map<long long, vector<int>>> vert_time;
map<long long, set<long long>> vert_by_time;
vector<vector<int>> g;

void dfs(long long v, long long t, long long e, long long ind, long long p) {
    path.push_back(v);
    if (v == e) {
        time_end[ind] = t + path.size() - 1;
        for (auto &i : path) {
            vert_by_time[t].insert(i);
            all_times.insert(t);
            vert_time[i][t].push_back(ind);
            t++;
        }
        path.pop_back();
        return;
    }
    // cout << v << ' ' << g.size() << endl;
    for (auto &i : g[v]) {
        if (i == p) {
            continue;
        }
        // cout << "EDGE " << i << " " << v << endl;
        dfs(i, t, e, ind, v);
    }
    path.pop_back();
}

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
) {
    long long n = R.size() + 1;
    long long m = D.size();
    g.resize(n);
    for (auto &[u, v] : R) {
        // cout << u << " " << v << endl;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    vector<long long> time_inf(m);
    time_end.resize(m);
    vert_time.resize(n);
    // cout << "HERE" << endl;
    for (long long i = 0; i < m; i++) {
        dfs(S[i], D[i], T[i], i, -1);
        if (I[i]) {
            time_inf[i] = D[i];
        } else {
            time_inf[i] = 1e18;
        }
    }
    for (auto &i : all_times) {
        for (auto &lll : vert_by_time[i]) {
            auto j = vert_time[lll];
                bool inf = false;
                for (auto &k : j[i]) {
                    if (I[k]) {
                        inf = true;
                    }
                }
                if (inf) {
                    for (auto &k : j[i]) {
                        I[k] = true;
                        time_inf[k] = min(time_inf[k], i);
                    }
                }
        }
    }
    vector<int> ans(m);
    for (long long i = 0; i < m; i++) {
        ans[i] = max(0ll, time_end[i] - time_inf[i] + 1);
    }
    return ans;
}