#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;
vector <int> depth;
int timer = 0;

using llint = long long;

void dfs (int p, int par, int d) {
    depth[p] = d;
    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, d + 1);
    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
) {
    const 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);
    depth.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, 1);

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

    vector <pair <int, int> > route_size (S.size());
    for (int i = 0; i < S.size(); i++) {
        int a = S[i], b = T[i];
        const int c = lca (a, b);
        route_size[i].first = i;
        route_size[i].second =  (
            depth[a] - depth[c] +
            depth[b] - depth[c] + 1
        );
    }
    sort (route_size.begin(), route_size.end(), 
        [](auto&& a, auto&& b) {
            return a.second < b.second;
        }
    );

    {
        auto D1 = D;
        for (int i = 0; i < D.size(); i++)
            D[route_size[i].first] = D1[i];
        auto S1 = S;
        for (int i = 0; i < S.size(); i++)
            S[route_size[i].first] = S1[i];
        auto T1 = T;
        for (int i = 0; i < T.size(); i++)
            T[route_size[i].first] = T1[i];
        auto I1 = I;
        for (int i = 0; i < I.size(); i++)
            I[route_size[i].first] = I1[i];
    }

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

        int ind = 0;
        while (a != c) {
            route[ind++] = a;
            a = parent[a];
        }
        route[ind++] = c;
        
        ind = route.size() - 1;        
        while (b != c) {
            route[ind--] = b;
            b = parent[b];
        }
    }


    int t = 0;
    vector <int> ans (S.size());
    for (int _ = 0; _ < N; _++) {
        vector <vector <int> > ids (N + 1);
        vector <int> tchk;
        tchk.reserve (N + 1);
        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]);
        }
        tchk.erase (unique (tchk.begin(), tchk.end()), tchk.end());
        for (auto&& i : tchk)
            for (auto&& e : ids[i])
                I[e] = 1;
        for (int e = 0; e < S.size(); e++)
            ans[e] += I[e];
        t++;
    }
    
    vector <int> actans (S.size());
    for (int i = 0; i < S.size(); i++)
        actans[i] = ans[route_size[i].first];
    return actans;
}
