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

using namespace std;
const int MAX_N = 15e4+42;

int par[MAX_N][20], depth[MAX_N];
vector<int> adj_list[MAX_N];

void dfs(int u, int p, int d){
    depth[u] = d;
    par[u][0] = p;
    for(int v : adj_list[u])
        if(v != p) // tree
            dfs(v, u, d + 1);
}

int lift_k_up(int v, int k){
    for(int i = 0; i < 20; ++i)
        if(k & (1 << i))
            v = par[v][i];
    return v;
}

int LCA(int a, int b){
    if(depth[a] < depth[b]) swap(a, b); // we need a to be deeper than be
    a = lift_k_up(a, depth[a] - depth[b]); // a and b are on the same depth
    for(int i = 19; i >= 0; --i){
        if(lift_k_up(a, (1 << i)) != lift_k_up(b, (1 << i))){
            a = lift_k_up(a, (1 << i));
            b = lift_k_up(b, (1 << i));
        }
    }
    return par[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
) {
    vector<int> ans;
    int m = D.size();
    for(int i = 0; i < R.size(); ++i){
        int u = R[i].first, v = R[i].second;
        adj_list[u].push_back(v);
        adj_list[v].push_back(u);
    }
    dfs(0, 0, 0); // does the LCA precomps
    for(int i = 0; i < m; ++i){
        if(I[i] == 0){
            ans.push_back(i);
            continue;
        }
        int m = LCA(S[i], T[i]);
        // length of the path is from start to LCA and from LCA to T plus one (kolcheta i dupki)
        ans.push_back(depth[m] - depth[S[i]] + depth[m] - depth[T[i]] + 1);
    }
    return ans;
}
