#include "infection.h"
#include <vector>
#include <utility>
using namespace std;
vector<vector<int>> par(150005,vector<int>(21));
vector<vector<int>> g(150005);
vector<int> depth(150005,0);
void dfs1(int node, int p, int d){
    depth[node]=d;
    par[node][0]=p;
    for(int nei:g[node])dfs1(nei,node,d+1);
}
int lca(int a, int b){
    if(a==b)return a;
    if(depth[a]<depth[b])swap(a,b);
        for(int i=20;i>=0;i--){
            if(depth[par[a][i]]>=depth[b])a=par[a][i];
        }
    for(int i=20;i>=0;i--){
            if(par[a][i]!=par[b][i])a=par[a][i],b=par[b][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
) {
    int n=R.size()+1;    int m=I.size();

    for(auto p:R){
        g[p.first].push_back(p.second);
        g[p.second].push_back(p.first);
    }
    dfs1(0,0,0);
    for(int i=1;i<=20;i++){
        for(int j=0;j<n;j++){
            par[j][i]=par[par[j][i-1]][i-1];
        }
    }
    vector<int> res(m);
    for(int i=0;i<m;i++){
        res[i]=I[i]*(depth[S[i]]+depth[T[i]]-2*depth[lca(S[i],T[i])]+1);
    }
    return res;
    
}
