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

using namespace std;

#define ff first
#define ss second
#define all(a) a.begin(),a.end()
#define pb push_back
#define endl '\n'

std::vector<int> solve(
    std::vector<std::pair<int, int> > edge,
    std::vector<long long> d,
    std::vector<int> s,
    std::vector<int> t,
    std::vector<bool> i1
) {


    int n = edge.size() + 1;
    vector <vector <int> > g(n + 1);



    for(int i = 0;i < n - 1;i++){
        g[edge[i].ff].pb(edge[i].ss);
        g[edge[i].ss].pb(edge[i].ff);
    }

    int m = d.size();

    vector <int> c(m);

    for(int i = 0;i < m;i++){
        if(i1[i]){
            c[i] = 1;
        }
    }

    vector <vector <int> > child(m , vector <int> (n + 1, -1));

    set <int> st;

    for(int i = 0;i < m;i++){
        vector <int> par(n + 1, -1);

        queue <int> q;

        q.push(s[i]);
        par[s[i]] = s[i];

        int den = d[i];

        while(!q.empty()){
            int x = q.front();
            q.pop();

            for(int u = 0;u < (int)g[x].size();u++){
                int e = g[x][u];
                if(par[e] == -1){
                    par[e] = x;
                    q.push(e);
                }
            }
        }
        while(par[t[i]] != t[i]){
            child[i][par[t[i]]] = t[i];
            t[i] = par[t[i]];
            st.insert(den);
            den++;
        }
    }
    vector <int> ans(m);

    for(int den : st){
        //cout<<den<<endl;
        vector <int> inf(n);
        for(int i = 0;i < m;i++){
            if(s[i] == -1) continue;
            if(d[i] == den && c[i] == 1){
                inf[s[i]] = 1;
            }
        }
        for(int i = 0;i < m;i++){
            if(den != d[i]) continue;
            if(s[i] == -1) continue;
            c[i] = max(c[i] , inf[s[i]]);
            ans[i] += c[i];
            s[i] = child[i][s[i]];
            d[i]++;
        }
    }

    return ans;

}
