#include "infection.h"
#include <bits/stdc++.h>

using namespace std;

#define all(x) (x).begin(), (x).end()

const int N = 15e4 + 5;
using ll = long long;

struct Ev {
	ll timp;
	int loc, persoana;
	bool operator < (const Ev &oth) const {
		return timp < oth.timp;
	}
};


vector<int> vec[N];
int depth[N];
int lift[20][N];
bool infectat[N];
ll last_infected[N]; //pentru oras sau ceva

int lca(int u, int v) {
	if (depth[u] > depth[v]) {
		swap(u, v);
	}
	for (int bit = 19; bit >= 0; bit --) {
		if (depth[v] - (1 << bit) >= depth[u]) {
			v = lift[bit][v];
		}
	}
	if (u == v) {
		return u;
	}
	for (int bit = 19; bit >= 0; bit --) {
		if (lift[bit][u] != lift[bit][v]) {
			v = lift[bit][v];
			u = lift[bit][u];
		}
	}
	return lift[0][v];
}

int dist(int u, int v) {
	int w = lca(u, v);
	return depth[u] + depth[v] - 2 * depth[w] + 1;
}

vector<int> path(int l, int r) {
	int m = lca(l, r);
	vector<int> ans;
	vector<int> b1, b2;
	while (l != m) {
		b1.push_back(l);
		l = lift[0][l];
	}
	while (r != m) {
		b2.push_back(r);
		r = lift[0][r];
	}
	reverse(b2.begin(), b2.end());
	for (auto i : b1) {
		ans.push_back(i);
	}
	ans.push_back(m);
	for (auto i : b2) {
		ans.push_back(i);
	}
	return ans;
}

void dfs(int nod, int papa) {
	depth[nod] = 1 + depth[papa];
	lift[0][nod] = papa;
	for (int bit = 1; bit < 20; bit ++) {
		lift[bit][nod] = lift[bit - 1][lift[bit - 1][nod]];
	}
	for (auto i : vec[nod]) {
		if (i == papa) {
			continue;
		}
		dfs(i, nod);
	}
}

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 (int i = 0; i < n - 1; i ++) {
		auto [u, v] = R[i];
		u ++, v ++;
		vec[u].push_back(v);
		vec[v].push_back(u);
	}
	for (int i = 1; i <= n; i ++) {
		last_infected[i] = -1;
	}
	vector<Ev> ev;
	dfs(1, 0);
	for (int i = 0; i < m; i ++) {
		int l = S[i] + 1, r = T[i] + 1;
		ll time = D[i];
		vector<int> pth = path(l, r);
		for (int j = 0; j < (int)pth.size(); j ++) {
			ev.push_back({time + j, pth[j], i});
		}
		infectat[i] = I[i];
	}
	sort(all(ev));
	vector<int> ans(m, 0);
	int pos = 0;
	while (pos < (int)ev.size()) {
		int lpos = pos + 1;
		while (lpos < (int)ev.size() && ev[lpos].timp == ev[pos].timp) {
			lpos ++;
		}
		for (int i = pos; i < lpos; i ++) {
			if (infectat[ev[i].persoana]) {
				last_infected[ev[i].loc] = ev[i].timp;
			}
		}
		for (int i = pos; i < lpos; i ++) {
			if (last_infected[ev[i].loc] == ev[i].timp) {
				infectat[ev[i].persoana] = 1;
				ans[ev[i].persoana] ++;
			}
		}
		pos = lpos;
	}
	return ans;
}
