#include "infection.h"

#include <algorithm>
#include <functional>
#include <ranges>
#include <unordered_set>
#include <unordered_map>
using namespace std;

vector<int> solve(vector<pair<int, int>> R, vector<long long> D, vector<int> S, vector<int> T, vector<bool> I) {
	const int n = R.size() + 1, m = D.size();
	vector<vector<int>> g(n), paths(m);
	for (auto &[u, v] : R) {
		g[u].push_back(v);
		g[v].push_back(u);
	}
	const function<bool(int, int, int, vector<int>&)> dfs = [&](const int u, const int p, const int f, vector<int> &vec) {
		if (u == f) return vec.push_back(u), true;
		for (const int v : g[u]) {
			if (v == p) continue;
			if (dfs(v, u, f, vec)) return vec.push_back(u), true;
		}
		return false;
	};
	vector<pair<long long, int>> activate(m);
#pragma GCC ivdep
	for (int i = 0; i < m; i++) {
		activate[i] = {D[i], i};
		dfs(S[i], -1, T[i], paths[i]);
	}
	ranges::sort(activate, greater());
	unordered_set<int> active;
	long long dayNow = 0;
	vector<int> ans(m);
	while (!activate.empty() || !active.empty()) {
		if (active.empty()) {
			dayNow = activate.back().first;
		}
		vector<int> forRemoval;
		for (const int i : active) {
			if (I[i]) ++ans[i];
			paths[i].pop_back();
			if (paths[i].empty()) forRemoval.push_back(i);
		}
		for (const int i : forRemoval) active.erase(i);
		while (!activate.empty() && activate.back().first == dayNow) {
			active.insert(activate.back().second);
			activate.pop_back();
		}
		unordered_map<int, vector<int>> parties;
		for (const int i : active) parties[paths[i].back()].push_back(i);
#pragma GCC ivdep
		for (auto &v: parties | views::values) {
			for (const int i : v) if (I[i]) {
#pragma GCC ivdep
				for (const int j : v) I[j] = true;
				break;
			}
		}
		++dayNow;
	}
	return ans;
}