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

#define rep(i, s, e) for (ll i = s; i < e; i++)
#define upmax(a, b) a = max(a, b)
#define upmin(a, b) a = min(a, b)

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>;

const ll INF = 1e18;
const ll MAX_JUMP = 25;

vvll tree, papa;
vll depth;

void dfs(ll node, ll p) {
    for (auto& nei : tree[node]) {
        if (nei == p) continue;
        depth[nei] = depth[node] + 1;
        dfs(nei, node);
        papa[0][nei] = node;
    }
}

ll LCA(ll a, ll b) {
    if (depth[b] > depth[a]) swap(a, b);
    for (ll i = MAX_JUMP - 1; i >= 0; i--) {
        if (papa[i][a] != -1 && depth[papa[i][a]] >= depth[b]) {
            a = papa[i][a];
        }
    }
    if (a == b) return a;
    for (ll i = MAX_JUMP - 1; i >= 0; i--) {
        if (papa[i][a] != papa[i][b]) {
            a = papa[i][a];
            b = papa[i][b];
        }
    }
    return papa[0][a];
}

vector<int> solve(vector<pair<int, int>> R, vector<long long> D, vector<int> S, vector<int> T, vector<bool> I) {
    ll n = R.size() + 1;
    ll m = T.size();
    tree.clear(), tree.resize(n);
    depth.clear(), depth.resize(n);
    papa.clear(), papa.resize(MAX_JUMP, vll(n, -1));
    vector<pair<pair<pll, pll>, ll>> arr(m);
    rep(i, 0, n - 1) {
        tree[R[i].first].push_back(R[i].second);
        tree[R[i].second].push_back(R[i].first);
    }
    rep(i, 0, m) {
        arr[i] = { { {I[i], D[i]}, {S[i], T[i]} }, i};
    }
    dfs(0, -1);
    rep(j, 1, MAX_JUMP) {
        rep(i, 0, n) {
            if (papa[j - 1][i] != -1) papa[j][i] = papa[j - 1][papa[j - 1][i]];
        }
    }
    sort(arr.rbegin(), arr.rend());
    vector<int> ans(m);
    rep(i, 0, m) {
        ll is_infected = arr[i].first.first.first;
        ll day = arr[i].first.first.second;
        ll a = arr[i].first.second.first;
        ll b = arr[i].first.second.second;
        ll idx = arr[i].second;
        ll lca = LCA(a, b);
        if (is_infected) ans[idx] = depth[a] + depth[b] - 2 * depth[lca] + 1;
    }
    return ans;
}
