#include "infection.h"
#include <vector>
#include <utility>
#include<iostream>
#include<algorithm>
#include<vector>
#include<string>
#include<set>
#include<map>
#include<queue>
#include<stack>
#include<iomanip>
#include<math.h>

using namespace std;
using ll = long long;
using vll = vector<ll>;
using vvll = vector<vll>;
using pll = pair<ll, ll>;
using vpll = vector<pll>;
using vvpll = vector<vpll>;

vvll tree;
vll father;
vll h;

void dfs(ll node, ll prev) {
    for (auto i : tree[node]) {
        if (i == prev) { continue; }
        h[i] = h[node] + 1;
        father[i] = node;
        dfs(i, node);
    }
}

ll max_jump = 20;
vvll lca;

ll find_lca(ll x, ll y) {
    if (h[x] < h[y]) { swap(x, y); }
    ll k = h[x] - h[y];
    for (ll i = max_jump; i >= 0; i--) {
        if (1LL << i <= k) {
            k -= 1LL << i;
            x = lca[i][x];
        }
    }
    if (x == y) { return x; }
    if (lca[0][x] == lca[0][y]) { return lca[0][x]; }
    for (ll i = 1; i <= max_jump; i++) {
        if (lca[i][x] == lca[i][y]) {
            return find_lca(lca[i - 1][x], lca[i - 1][y]);
        }
    }
}

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
) {
    ll n = R.size() + 1;
    tree.clear();
    h.clear();
    father.clear();
    father.resize(n);
    h.resize(n);
    tree.resize(n);
    for (ll i = 0; i < n - 1; i++) {
        tree[R[i].first].push_back(R[i].second);
        tree[R[i].second].push_back(R[i].first);
    }
    dfs(0, 0);
    ll m = D.size();
    lca.clear();
    lca.resize(max_jump + 1, vll(n));
    for (ll i = 0; i < n; i++) {
        lca[0][i] = father[i];
    }
    for (ll i = 1; i <= max_jump; i++) {
        for (ll j = 0; j < n; j++) {
            lca[i][j] = lca[i - 1][lca[i - 1][j]];
        }
    }
    vector<int> ans(m);
    vll papa(m);
    for (ll i = 0; i < m; i++) {
        papa[i] = find_lca(S[i], T[i]);
    }
    for (ll i = 0; i < m; i++) {
        ans[i] = I[i] * (h[S[i]] + h[T[i]] - 2*h[papa[i]] + 1);
    }
    return ans;
    
}
