#include "infection.h"
#include <vector>
#include <utility>
#include <bits/stdc++.h>
#define ll long long
using namespace std;
vector<vector<int>> v, ord, up;
vector<int> pa, de;
int go = 0, ind = 0;
void dfs(int x, int p, int dd){
    pa[x] = p;
    de[x] = dd;
    up[x][0] = p;
    for (int i = 1; i < 20; i++){
        up[x][i] = up[up[x][i - 1]][i - 1];
    }
    for (auto el : v[x]){
        if (el != p){
            dfs(el, x, dd + 1);
        }
    }
}
int dist(int x, int k){
    for (int i = 0; i < 20; i++){
        int val = (k & (1<<i));
        if (val != 0) x = up[x][i];
    }
    return x;
}
int lca(int a, int b){
    if (de[a] > de[b]) swap(a, b);
    b = dist(b, de[b] - de[a]);
    if (a == b) return a;
    for (int i = 19; i >= 0; i--){
        if (up[a][i] != up[b][i]){
            a = up[a][i], b = up[b][i];
        }
    }
    return up[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, m = S.size();
    v.assign(n + 1, vector<int>()), ord.assign(m + 1, vector<int>()), pa.assign(n + 1, 0);
    up.assign(n + 1, vector<int>(20, 0)), de.assign(n + 1, 0);
    for (auto el : R){
        v[el.first].push_back(el.second);
        v[el.second].push_back(el.first);
    }
    dfs(0, 0, 0);
    vector<int> ans(m + 1, 0);
    for (int i = 0; i < m; i++){
        if (I[i] == 0) continue;
        int lc = lca(S[i], T[i]);
        ans[i] = de[S[i]] + de[T[i]] - 2 * de[lc] + 1;
    }
    return ans;
}
