#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 INTINF = 1e9 + 10;

int n, m, k;
std::vector<int> g[MAXN];
std::vector<int> ans;
int parent[MAXN];

void initDFS(int node, int par)
{
    parent[node] = par;

    for(const int &to : g[node])
    {
        if(to == par)
        {
            continue;
        }

        initDFS(to, node);
    }
}

std::vector<int> getPath(int start, int end)
{
    initDFS(start, -1);

    std::vector<int> path;

    int current = end;
    while(current != start)
    {
        path.push_back(current);
        current = parent[current];
    }

    path.push_back(start);
    std::reverse(path.begin(), path.end());
    return path;
}

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].push_back(v);
        g[v].push_back(u);
    }

    ans.resize(k, 0);
    for(int i = 0 ; i < k ; ++i)
    {
        if(I[i])
        {
            std::vector<int> path = getPath(S[i], T[i]);
            ans[i] = path.size();
        }
    }

    return ans;
}
