#include "infection.h"
#include <bits/stdc++.h>
using namespace std;

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 m = I.size(), n = R.size() + 1;
	vector<vector<int>> p(n, vector<int>(20)), g(n);
	vector<int> tin(n), tout(n), dep(n);
	int timer = 0;
	
	for(auto to : R){
		g[to.first].push_back(to.second);
		g[to.second].push_back(to.first);
	}
	
	auto dfs = [&](auto self, int x, int pr) -> void {
		p[x][0] = pr;
		for(int i = 1;i<20;i++)p[x][i] = p[p[x][i-1]][i-1];
		tin[x] = timer++;
		for(auto to : g[x]){
			if(to == pr)continue;
			dep[to] = dep[x] + 1;
			self(self, to, x);
		}
		tout[x] = timer++;
	};
	
	dfs(dfs, 0, 0);
	
	auto is = [&](int a, int b) -> bool {
		return tin[a] <= tin[b] && tout[b] <= tout[a];
	};
	
	auto lca = [&](int a, int b) -> int {
		if(is(a, b))return a;
		if(is(b, a))return b;
		for(int i = 19;i>=0;i--){
			if(!is(p[a][i], b)){
				a = p[a][i];
			}
		}
		
		a = p[a][0];
		return a;
	};
	
	vector<int>res(m);
	for(int i = 0;i<m;i++){
		if(!I[i])continue;
		int u = lca(S[i], T[i]);
		res[i] = dep[S[i]] + dep[T[i]] - dep[u] * 2 + 1;
	}
	
	return res;
		
	
}
