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

using namespace std;

constexpr int MAXN = 150000;
constexpr int MAXM = 150000;

#define sz(x) ((int)((x).size()))

using ll = long long;

struct Person
{
	int d, s, t;
	bool z;
	Person(int d, int s, int t, bool z) : d(d), s(s), t(t), z(z) {}
	Person() : d(), s(), t(), z() {}
};

int n, m, root;
Person v[MAXM+2];
int ans[MAXM+2];
vector<int> nb[MAXN+2];
int p[MAXN+2], sz[MAXN+2], depth[MAXN+2], jump[MAXN+2];

void dfs(int node)
{
	int bigjump = (depth[node] - depth[jump[node]] == depth[jump[node]] - depth[jump[jump[node]]]) ? jump[jump[node]] : node;
	sz[node] = 1;
	for (auto x : nb[node]) {
		if (!p[x]) {
			p[x] = node;
			depth[x] = depth[node]+1;
			jump[x] = bigjump;
			dfs(x);
			sz[node] += sz[x];
		}
	}
}

int getLca(int a, int b)
{
	if (depth[a] < depth[b]) swap(a, b);
	while (depth[a] > depth[b]) {
		if (depth[jump[a]] >= depth[b]) {
			a = jump[a];
		}
		else {
			a = p[a];
		}
	}
	while (a != b) {
		if (jump[a] != jump[b]) {
			a = jump[a], b = jump[b];
		}
		else {
			a = p[a], b = p[b];
		}
	}
	return a;
}

vector<int> getChain(int a, int b)
{
	vector<int> lant;
	int lca = getLca(a, b);
	while (a != lca) {
		lant.push_back(a);
		a = p[a];
	}
	lant.push_back(lca);
	int poz = sz(lant);
	while (b != lca) {
		lant.push_back(b);
		b = p[b];
	}
	reverse(lant.begin()+poz, lant.end());
	return lant;
}

vector<int> bruteSolve()
{
	vector<pair<pair<int, int>, int>> evs;
	for (int i = 1; i <= m; i++) {
		vector<int> lant = getChain(v[i].s, v[i].t);
		for (int j = 0; j < sz(lant); j++) {
			evs.push_back({{v[i].d+j, lant[j]}, i});
		}
		lant.clear();
	}
	sort(evs.begin(), evs.end());
	for (int i = 0; i < sz(evs); i++) {
		int j = i;
		while (j < sz(evs)-1 && evs[j+1].first == evs[i].first) j++;
		bool ok = false;
		for (int aux = i; aux <= j; aux++) {
			if (v[evs[aux].second].z) {
				ok = true;
				break;
			}
		}
		if (ok) {
			for (int aux = i; aux <= j; aux++) {
				v[evs[aux].second].z = true;
			}
		}
		for (int aux = i; aux <= j; aux++) {
			ans[evs[aux].second] += v[evs[aux].second].z;
		}
		i = j;
	}
	vector<int> ansvec;
	for (int i = 1; i <= m; i++) {
		ansvec.push_back(ans[i]);
	}
	return ansvec;
}

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
) {
    n = sz(R)+1;
    for (int i = 0; i < n; i++) {
		nb[R[i].first+1].push_back(R[i].second+1);
		nb[R[i].second+1].push_back(R[i].first+1);
	}
	m = sz(D);
	for (int i = 0; i < m; i++) {
		v[i+1] = Person(D[i], S[i]+1, T[i]+1, I[i]);
	}
	root = 1;
	p[root] = root;
	depth[root] = 1;
	jump[root] = root;
	dfs(root);
	if (n <= 5000 && m <= 5000) {
		return bruteSolve();
	}
	return {};
}
