#include "infection.h"
#include <vector>
#include <utility>
#include <algorithm>
#include <set>

using namespace std;

const int N=5010;

vector<int> tr[N],path[N];
vector<int> S,T;
int n,m,cnt[N];

bool dfs(int u,int fa,int id)
{
    if(u==T[id])
    {
        path[id].push_back(u);
        return true;
    }
    for(int v:tr[u])
    {
        if(v==fa) continue;
        if(dfs(v,u,id))
        {
            path[id].push_back(u);
            return true;
        }
    }
    return false;
}

std::vector<int> solve(
    std::vector<std::pair<int, int>> R,
    std::vector<long long> D,
    std::vector<int> SS,
    std::vector<int> TT,
    std::vector<bool> I
) {
    S=SS,T=TT;
    n=R.size()+1,m=S.size();
    for(int i=0;i<n;i++) tr[i].clear();
    for(int i=0;i<n;i++) cnt[i]=0;
    for(int i=0;i<n-1;i++)
    {
        auto [u,v]=R[i];
        tr[u].push_back(v);
        tr[v].push_back(u);
    }
    for(int i=0;i<m;i++) path[i].clear();
    for(int i=0;i<m;i++)
    {
        dfs(S[i],-1,i);
    }
    vector<int> seq(m);
    set<int> s;
    for(int i=0;i<m;i++) seq[i]=i;
    sort(seq.begin(),seq.end(),[D](int i,int j){return D[i]<D[j];});

    long long time=-1,now=0;
    vector<int> ans(m,0);
    while(!s.empty()||now<=m-1)
    {
        time++;
        auto it=s.begin();
        while(it!=s.end())
        {
            int k=*it;
            int v=path[k][path[k].size()-1];
            path[k].pop_back();
            cnt[v]-=I[k];
            if(path[k].size()==0)
            {
                auto it2=it;
                it++;
                s.erase(it2);
            }
            else
            {
                int u=path[k][path[k].size()-1];
                cnt[u]+=I[k];
                it++;
            }
        }
        if(s.empty())
        {
            time=D[seq[now]];
        }
        while(now<=m-1&&D[seq[now]]==time)
        {
            int u=path[seq[now]][path[seq[now]].size()-1];
            cnt[u]+=I[seq[now]];
            s.insert(seq[now]);
            now++;
        }
        for(auto k:s)
        {
            int u=path[k][path[k].size()-1];
            if(cnt[u]) ans[k]=max(ans[k],(int)path[k].size());
        }
    }

    return ans;
}
