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