#include "infection.h"
#include<iostream>
#include <vector>
#include <utility>
#include <algorithm>
#define MAXN 5010
using namespace std;
long long day[MAXN];
int st[MAXN];
int en[MAXN];
long long sick[MAXN];
vector<int> edges[MAXN];
vector<long long> order[MAXN];//each node, each person at what day
bool dfs(int curr, int parent, long long dep, const int init)
{
    if(curr==en[init]){order[curr][init]=dep;return 1;}
    for(int i=0;i<edges[curr].size();i++)
    {
        int now=edges[curr][i];
        if(now==parent)continue;
        if(dfs(now,curr,dep+1,init))
        {
            order[curr][init]=dep;
            return 1;
        }
    }
    return 0;
}
vector<int> solve(vector<pair<int, int>> R,vector<long long> D,vector<int> S,vector<int> T,vector<bool> I)
{
    for(int i=0;i<R.size();i++)
    {
        edges[R[i].first].push_back(R[i].second);
        edges[R[i].second].push_back(R[i].first);
    }
    int n=R.size()+1;
    int m=D.size();
    for(int i=0;i<n;i++)
    {
        order[i].resize(m);
        for(int j=0;j<m;j++)order[i][j]=-1;
    }
    for(int i=0;i<m;i++)
    {
        day[i]=D[i];
        st[i]=S[i];
        en[i]=T[i];
        if(I[i])sick[i]=D[i];
        else sick[i]=LLONG_MAX;
    }
    vector<pair<long long, int>> sorted;
    for(int i=0;i<m;i++)//get paths
    {
        dfs(st[i],-1,day[i],i);
        sorted.push_back({day[i],i});
    }
    sort(sorted.begin(),sorted.end());

    /*for(int i=0;i<n;i++)
    {
        for(int j=0;j<m;j++)cout<<order[i][j]<<" ";
        cout<<"A\n";
    }*/

    vector<int> ans(m);
    for(int z=0;z<sorted.size();z++)
    {
        int i=sorted[z].second;
        long long val=sorted[z].first;
        //cout<<i<<" "<<val<<"Begin\n";
        int curr=st[i],cntr=0;
        while(curr!=en[i])
        {
            int nxt=0;
            for(int f=0;f<n;f++)
            {
                if(order[f][i]==val+1){nxt=f;break;}
            }
            if(sick[i]<=val){cntr++;curr=nxt;val++;continue;}
            for(int j=0;j<m;j++)
            {
                if(order[curr][j]==order[curr][i] && order[curr][j]!=-1)
                {
                    if(sick[j]<=order[curr][i])
                    {
                        //cout<<j<<" infected i at day "<<order[curr][i]<<"\n";
                        sick[i]=order[curr][i];
                        cntr++;
                        break;
                    }
                }
            }
            curr=nxt;
            val++;
        }
        if(sick[i]<=val)cntr++;
        else
        {
            for(int j=0;j<m;j++)
            {
                if(order[curr][j]==order[curr][i] && order[curr][j]!=-1)
                {
                    if(sick[j]<=order[curr][i]){sick[i]=order[curr][i];cntr++;break;}
                }
            }
        }
        ans[i]=cntr;
    }
    return ans;
}