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

vector <vector <int> > gr;
vector <vector <int> > binlift;
vector <int> in;
vector <int> out;
vector <int> parent;
int timer = 0;

using llint = long long;

void dfs (int p, int par) {
    parent[p] = par;
    in[p] = out[p] = ++timer;

    binlift[p][0] = par;
    for (int i = 1; i < 20; i++)
        binlift[p][i] = binlift[binlift[p][i - 1]][i - 1];

    for (int& e : gr[p])
        if (e != par)
            dfs (e, p);
    out[p] = timer;
}

inline bool is_ancestor (int a, int b) {
    return in[a] <= in[b] && in[b] <= out[a];
}

int lca (int a, int b) {
    if (is_ancestor (a, b))
        return a;
    if (is_ancestor (b, a))
        return b;
    for (int i = 19; i >= 0; i--)
        if (!is_ancestor (binlift[a][i], b))
            a = binlift[a][i];
    return binlift[a][0];
}

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
) {
    int N = R.size() + 1;
    gr.resize (N + 1);
    binlift.resize (N + 1, vector <int> (20));
    in.resize (N + 1);
    out.resize (N + 1);
    parent.resize (N + 1);
    out[0] = 2e9;
    for (auto [a, b] : R) {
        a++, b++;
        gr[a].emplace_back (b);
        gr[b].emplace_back (a);
    }

    dfs (1, 0);

    for (int& e : S)
        e++;
    for (int& e : T)
        e++;

    vector <vector <int> > routes (S.size());
    for (int i = 0; i < S.size(); i++) {
        int a = S[i], b = T[i];
        int c = lca (a, b);
        vector <int>& route = routes[i];

        while (a != c) {
            route.emplace_back (a);
            a = parent[a];
        }
        route.emplace_back (c);
        
        vector <int> temp;
        while (b != c) {
            temp.emplace_back (b);
            b = parent[b];
        }
        reverse (temp.begin(), temp.end());
        for (int& e : temp)
            route.emplace_back (e);
    }

    int t = 0;
    vector <int> ans (S.size());
    for (int _ = 0; _ < N; _++) {
        vector <vector <int> > ids (N + 1);
        vector <int> tchk;
        for (int e = 0; e < S.size(); e++) {
            if (t >= routes[e].size()) {
                I[e] = 0;
                continue;
            }
            ids[routes[e][t]].emplace_back (e);
            if (I[e])
                tchk.emplace_back (routes[e][t]);
        }
        for (auto&& i : tchk)
            for (auto&& e : ids[i])
                I[e] = 1;
        for (int e = 0; e < S.size(); e++)
            if (I[e])
                ans[e]++;
        t++;
    }
    
    return ans;
}
