#include "infection.h"
#include <vector>
#include <utility>
#include <algorithm>

using namespace std;
typedef long long ll;

const int N = 150000 + 7;

vector<int> g[N];
int sz[N], p[N];

void hld_reorder(int v, int pr) {
    sz[v] = 1;
    p[v] = pr;

    for (int i = 0; i < g[v].size(); ++i) {
        if (g[v][i] == p[v]) {
            g[v].erase(g[v].begin() + i);
            break;
        }
    }

    for (int u : g[v]) {
        hld_reorder(u, v);
        sz[v] += sz[u];
    }

    sort(g[v].begin(), g[v].end(), [&](int x, int y) {
        return sz[x] > sz[y];
    });
}

int tim = 0;
int tin[N], tout[N], h[N], st[N];
void dfs(int v, int s) {
    st[v] = s;
    tin[v] = tim++;
    h[v] = h[p[v]] + 1;

    if (!g[v].empty()) {
        dfs(g[v][0], s);
    }
    for (int i = 1; i < g[v].size(); ++i) {
        dfs(g[v][i], g[v][i]);
    }
    tout[v] = tim;
}

int get(int x, int y){
    int ans = 0;
    while (st[x] != st[y]) {
        if (h[x] > h[y]) {
            ans += h[x];
            x = p[st[x]];
            ans -= h[x];
        }
        else {
            ans += h[y];
            y = p[st[y]];
            ans -= h[y];
        }
    }
    ans += abs(h[x] - h[y]);
    return ans;
}

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
) 
{
    for (int i = 0; i < R.size(); ++i) {
        g[R[i].first].push_back(R[i].second);
        g[R[i].second].push_back(R[i].first);
    }

    hld_reorder(0, 0);
    dfs(0, 0);

    int q = S.size();
    
    vector<int> ans(q, 0);
    for (int i = 0; i < q; ++i) {
        if (I[i]) {
            ans[i] = get(S[i], T[i]) + 1;
        }
    }
    return ans;
}
