#include<bits/stdc++.h>
#include "infection.h"
#define ll long long
#define ss second
#define ff first
#define control std::cout<<"here"<<'\n';
const int maxn = 5 * 1e3 + 5;
std::vector<int>v[maxn];
bool inf[maxn];
std::vector<int> d[maxn][maxn];
ll De[maxn], Se[maxn];
bool dfs(int node, int par, int depth, int st, int kr)
{
    if(node == kr)
    {
        d[De[node] + depth][node].push_back(st);
        return true;
    }
    int sz = v[node].size();
    for(int i = 0; i < sz; i ++)
    {
        int nb = v[node][i];
        if(nb != par)
        {
            bool b = dfs(nb, node, depth + 1, st, kr);
            if(b)
            {
                d[De[node] + depth][node].push_back(st);
                return true;
            }
        }
    }
    return false;
}
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)
{

    for(int i = 0 ; i < D.size(); i ++)
    {
        De[i] = D[i];
        Se[i] = S[i];
        inf[i] = I[i];
        //std::cout << inf[i] << " ";
    }
   // std::cout << '\n';
    for(int i = 0; i < R.size(); i ++)
    {
        v[R[i].ff].push_back(R[i].ss);
        v[R[i].ss].push_back(R[i].ff);
    }
    int n = R.size() + 1;
    int m = D.size();
    std::vector<int>ans(m);
    for(int i = 0; i < m; i ++)
    {
        dfs(S[i], -1, 1, i, T[i]);
    }
    for(int i = 1; i <= n + 2; i ++)
    {
        for(int j = 0; j < n; j ++)
        {
            bool infec = false;
            int infected;
            for(int z = 0 ; z < d[i][j].size(); z ++)
            {
                if(inf[d[i][j][z]])
                {infec = true;
                infected = d[i][j][z];}

            }
            if(infec)
            {
                for(int z = 0; z < d[i][j].size(); z ++)
                {
                    inf[d[i][j][z]] = true;
                    //if(d[i][j][z] == 1) std::cout << i << " " << j << " " << z << " " << infected << '\n';
                    ans[d[i][j][z]] ++;
                }
            }
        }
    }
    return ans;
}
