#include "infection.h"

#include <bits/stdc++.h>

using namespace std;

using i64 = long long;

vector<int> solve(vector<std::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>> adj(n), fa(n, vector <int> (n));

    for ( auto &[u, v]: R ){
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    for ( int x = 0; x < n; ++x ){
        auto dfs = [&](auto dfs, int u, int p) -> void{
            fa[x][u] = p;

            for ( auto &v: adj[u] ){
                if ( v != p ) dfs(dfs, v, u);
            }
        };

        dfs(dfs, x, -1);
    }

    map <i64, vector <pair<int, int>>> mp;

    vector <int> ans(m);
    vector <i64> f(m);

    for ( int i = 0; i < m; i++ ){
        int x = S[i], j = 0;

        while ( x != T[i] ){
            mp[D[i] + j].push_back({x, i});
            x = fa[T[i]][x], j += 1;
        }

        mp[D[i] + j].push_back({x, i});

        f[i] = D[i] + j + 1;
    }

    for ( auto &[it, b]: mp ){
        vector <bool> is(n);

        for ( auto &[v, i]: b ){
            if ( I[i] ) is[v] = 1;
        }

        for ( auto &[v, i]: b ){
            if ( !is[v] || ans[i] ) continue;

            ans[i] = f[i] - it;
            I[i] = 1;
        }
    }

    return ans;
}

