#include "infection.h"
#include <algorithm>
#include <iostream>
#include <cassert>
#include <numeric>
#include <random>
#include <stack>
#include <queue>
#include <set>
#include <map>

const int MAXN = 150000 + 10;
const int MAXLOG = 20;
const int INTINF = 1e9 + 10;

int n, m, k;
std::vector<int> ans;
std::vector<int> g[MAXN];
int parent[MAXN][MAXLOG];
int depth[MAXN];
int in[MAXN];
int out[MAXN];
int timer;

bool isAnc(int u, int v)
{
    return in[u] <= in[v] && out[u] >= out[v];
}

int getLCA(int u, int v)
{
    if(isAnc(u, v))
    {
        return u;
    }

    if(isAnc(v, u))
    {
        return v;
    }

    for(int lg = MAXLOG - 1 ; lg >= 0 ; --lg)
    {
        int next = parent[u][lg];
        if(next != 0 && !isAnc(next, v))
        {
            u = next;
        }
    }

    return parent[u][0];
}

void initDFS(int node, int par, int dp)
{
    depth[node] = dp;
    in[node] = ++timer;

    for(int lg = 0 ; lg < MAXLOG ; ++lg)
    {
        if(lg == 0)
        {
            parent[node][lg] = par;
            continue;
        }

        int prev = parent[node][lg - 1];
        parent[node][lg] = parent[prev][lg - 1];
    }
    
    for(const int &to : g[node])
    {
        if(to == par)
        {
            continue;
        }

        initDFS(to, node, dp + 1);
    }

    out[node] = timer;
}

int getPath(int start, int end)
{
    int lca = getLCA(start, end);
    //std::cout << lca << "\n";
    //std::cout << depth[start] << " " << depth[end] << " " << depth[lca] << "\n";
    return depth[start] + depth[end] - depth[lca] - depth[parent[lca][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
)
{
    n = R.size() + 1;
    k = D.size();

    for(auto &[u, v] : R)
    {
        g[u + 1].push_back(v + 1);
        g[v + 1].push_back(u + 1);
    }

    initDFS(1, 0, 1);

    ans.resize(k, 0);
    for(int i = 0 ; i < k ; ++i)
    {
        if(I[i])
        {
            ans[i] = getPath(S[i] + 1, T[i] + 1);
        }
    }

    return ans;
}
