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

using namespace std;

using pii = pair<int, int>;

bool dfs(int id, int v, int p, int goal, int dep, vector<vector<int>> &g, 
        map<pii, vector<int>> &e) {
    if (v == goal) {
        e[{dep, v}].push_back(id);
        return true;
    }
    for (int u : g[v]) {
        if (u != p) {
            if (dfs(id, u, v, goal, dep + 1, g, e)) {
                e[{dep, v}].push_back(id);
                return true;
            }
        }
    }
    return false;
}

vector<int> solve(vector<pair<int, int>> R, vector<long long> D, vector<int> S,
        vector<int> T, vector<bool> I) {
    int n = R.size() + 1, m = D.size();
    vector<vector<int>> g(n);
    for (auto [x, y] : R) {
        g[x].push_back(y);
        g[y].push_back(x);
    }

    map<pii, vector<int>> e;
    vector<int> ans(m);
    for (int i = 0; i < m; i++) {
        dfs(i, S[i], S[i], T[i], D[i], g, e);
    }

    /*for (auto [k, v] : e) {
        cout << k.first << " " << k.second << ": ";
        for (int x : v) cout << x << " ";
        cout << "\n";
    }*/

    for (auto [k, v] : e) {
        bool hasi = 0;
        for (int x : v) {
            hasi = hasi || I[x];
        }
        if (hasi) {
            for (int x : v) {
                I[x] = 1;
            }
        }
        for (int x : v) {
            ans[x] += I[x];
        }
    }

    return ans;
}
