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

vector<int> path;
vector<long long> time_end;
vector<long long> all_times;
vector<vector<int>> were;
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) {
        were[ind] = path;
        time_end[ind] = t + path.size() - 1;
        for (auto &i : path) {
            all_times.push_back(t);
            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);
    were.resize(n);
    // cout << "HERE" << endl;
    for (long long i = 0; i < m; i++) {
        dfs(S[i], D[i], T[i], i, -1);
    }

    sort(all_times.begin(), all_times.end());
    all_times.resize(unique(all_times.begin(), all_times.end()) - all_times.begin());
    vector<vector<long long>> by_time_all_v;
    vector<vector<pair<long long, long long>>> all(all_times.size());
    for (long long i = 0; i < m; i++) {
        long long t = D[i];
        t = lower_bound(all_times.begin(), all_times.end(), t) - all_times.begin();
        time_end[i] = lower_bound(all_times.begin(), all_times.end(), time_end[i]) - all_times.begin();
        if (I[i]) {
            time_inf[i] = t;
        } else {
            time_inf[i] = 1e18;
        }
        for (auto &j : were[i]) {
            all[t].push_back({j, i});
            t++;
        }
    }
    for (long long i = 0; i < all.size(); i++) {
        sort(all[i].begin(), all[i].end());
        vector<pair<long long, long long>> now;
        for (auto &j : all[i]) {
            // cout << i << " " << j.first << " " << j.second << endl;
            if (now.size() == 0 || j.first == now.back().first) {
                now.push_back(j);
            } else {
                //we have our time
                bool inf = false;
                for (auto &[tr, v] : now) {
                    if (I[v]) {
                        inf = true;
                    }
                } 
                if (inf) {
                    for (auto &[tr, v] : now) {
                        if (!I[v]) {
                            I[v] = true;
                            time_inf[v] = i;
                        }
                    }
                }
                now.clear();
                now.push_back(j);
            }
        }
        bool inf = false;
        for (auto &[tr, v] : now) {
            if (I[v]) {
                inf = true;
            }
        } 
        if (inf) {
            for (auto &[tr, v] : now) {
                if (!I[v]) {
                    I[v] = true;
                    time_inf[v] = i;
                }
            }
        }
        now.clear();
    }


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