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

using namespace std;

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;
    vector<int> g[n];
    for (auto [u, v] : R) {
        g[u].push_back(v);
        g[v].push_back(u);
    }
    vector<vector<int> > next(n, vector<int>(n, -1));
    function<void(int,int)> dfs = [&](int v, int o) {
        for (int to : g[v]) {
            if (next[to][o] == -1 and to != o) next[to][o] = v, dfs(to, o);
        }
    };
    for (int i = 0; i < n; i++) {
        dfs(i, i);
    }
    int m = D.size();
    auto cmp = [&](int i, int j) {
        return D[i] < D[j];
    };
    vector<int> ind(m);
    iota(ind.begin(), ind.end(), 0);
    sort(ind.begin(), ind.end(), cmp);
    vector<tuple<int,int,int,int> > X;
    long long now = 0, pt = 0;
    vector<int> ans(m);
    vector<tuple<int,int,int,int> > tmp;
    bool inf[n];
    while (true) {
        tmp.clear();
        for (auto [u, v, i, idx] : X) {
            if (u != v) tmp.push_back({next[u][v], v, i, idx});
        }
        X = tmp;
        
        if (X.empty()) {
            if (pt < m) now = D[ind[pt]];
            else break;
        }

        while (pt < m and D[ind[pt]] == now) X.push_back({S[ind[pt]], T[ind[pt]], I[ind[pt]], ind[pt]}), pt++;

        memset(inf, 0, sizeof(inf));
        for (auto [u, v, i, idx] : X) if (i) inf[u] = 1;
        for (auto &[u, v, i, idx] : X) {
            if (inf[u]) i = 1;
            ans[idx] += i;
        }
        now++;
    }
    return ans;
}
