#include "infection.h"
#include <vector>
#include <utility>
using namespace std;
const int NMAX=2e5+505;
int t[NMAX][20],lv[NMAX];
vector <int> v[NMAX];
void dfs(int nod,int ta)
{
    t[nod][0]=ta;
    lv[nod]=lv[ta]+1;
    for (int i=0;i<v[nod].size();++i)
    {
        if (v[nod][i]==t[nod][0]) continue;
        dfs(v[nod][i],nod);
    }
}
int lca(int x,int y)
{
    if (x==y) return x;
    if (lv[x]<lv[y]) swap(x,y);
    for (int i=18;i>=0;--i) {if (lv[y]+(1<<i)<lv[x]) {x=t[x][i];}}
    x=t[x][0];
    if (x==y) return x;
    for (int i=18;i>=0;--i) {if (t[x][i]==t[y][i]); else {x=t[x][i]; y=t[y][i];}}
    return t[x][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=n-1;
    for (int i=0;i<m;++i)
    {
        v[R[i].first].push_back(R[i].second);
        v[R[i].second].push_back(R[i].first);
    }
    lv[0]=0;
    dfs(0,0);
    for (int j=1;j<=18;++j)
    for (int i=0;i<n;++i)
    {
        t[i][j]=t[t[i][j-1]][j-1];
    }
    vector <int> ras;
    for (int i=0;i<n;++i)
    {
        if (I[i]==0) ras.push_back(0);
        else ras.push_back(lv[S[i]]+lv[T[i]]-2*lv[lca(S[i],T[i])]);
    }
    return ras;
}
