#include "infection.h"
#include <vector>
#include <utility>

using namespace std;

vector<int> adj[150005];
int tin[150005];
int tout[150005];
int d[150005];
int timer = 0;
int up[150005][22];

void dfs(int x,int p=-1) {
    
    tin[x] = timer;
    timer++;
    if (p != -1) d[x] = d[p] + 1;
    else d[p] = 1;

    up[x][0] = p;
    for(int k=1;k<22;++k) {
        if (up[x][k-1] == -1) up[x][k] = -1;
        else up[x][k] = up[ up[x][k-1] ][k-1];
    }

    for(auto e:adj[x]) {
        if (e == p) continue;

        dfs(e,x);
    }

    tout[x] = timer;

    return;
}

bool anc(int a,int b) {
    return (tin[a] < tin[b] and tout[a] >= tout[b]);
}

int lca(int a,int b) {
    if (a == b) return a;
    if (anc(a,b)) return a;
    if (anc(b,a)) return b;

    for(int k=21;k>=0;k--) {
        if (up[a][k] != -1 and !anc(up[a][k],b)) {
            a = up[a][k];
        }
    }

    return up[a][0];
}

int path(int S,int T) {
    int l = lca(S,T);
    return d[S] + d[T] - d[l]*2;
}


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 = (int)R.size() + 1;

    for(auto e:R) {
        adj[e.first].push_back(e.second);
        adj[e.second].push_back(e.first);
    }

    int M = (int)D.size();

    dfs(1);

    vector<int> res;

    for(int i=0;i<M;++i) {
        if (I[i] == 1) res.push_back(path(S[i],T[i])+1);
        else res.push_back(0);
    }

    return res;
}
