#include "infection.h"
#include "bits/stdc++.h"
using namespace std;

#define ll long long
const ll N=5001;
const ll INF=1e18;
vector<int>G[N];

ll dd[N][N],cc[N];
bool inf[N];
map<int,vector<pair<int,int>>>mp;

int go(int a,int b){
      ll mn=INF;
      for(int u:G[a]) mn=min(mn,dd[u][b]);
      for(int u:G[a]){
            if(dd[u][b]==mn) return u;
      }
}

vector<int>solve(vector<pair<int, int>>R,vector<long long>D,vector<int>S,vector<int>T,vector<bool>I)
{
      int n=(int)R.size()+1;
      for(int i=0;i<n-1;i++){
            G[R[i].first].push_back(R[i].second);
            G[R[i].second].push_back(R[i].first);
      }
      for(int i=0;i<n;i++){
            for(int j=0;j<n;j++) dd[i][j]=INF;
            queue<int>q;
            dd[i][i]=0;
            q.push(i);
            while(!q.empty()){
                  int v=q.front();
                  q.pop();
                  for(int u:G[v]){
                        if(dd[i][v]+1<dd[i][u]){
                              dd[i][u]=dd[i][v]+1;
                              q.push(u);
                        }
                  }
            }
      }
      int m=(int)S.size();
      set<ll>all;
      for(int i=0;i<m;i++){
            ll day=D[i];
            ll cur=S[i];
            ll fin=T[i];
            all.insert(day);
            mp[day].push_back({i,cur});
            while(1){
                  day++;
                  all.insert(day);
                  if(cur==fin) break;
                  cur=go(cur,fin);
                  mp[day].push_back({i,cur});
            }
      }
      vector<ll>ds;
      for(ll x:all) ds.push_back(x);
      vector<int>ans(m,0);
      for(ll dy:ds){
            for(pair<int,int>p:mp[dy]){
                  inf[p.second]=0;
            }
            for(pair<int,int>p:mp[dy]){
                  if(I[p.first]) inf[p.second]=1;
            }
            for(pair<int,int>p:mp[dy]){
                  if(inf[p.second]) I[p.first]=1;
                  ans[p.first]+=I[p.first];
            }
      }
      return ans;
}
