#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];

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();
      vector<vector<int>>all;
      for(int i=0;i<m;i++){
            ll day=D[i];
            ll cur=S[i];
            ll fin=T[i];
            all.push_back({day,i,cur});
            while(1){
                  day++;
                  if(cur==fin) break;
                  cur=go(cur,fin);
                  all.push_back({day,i,cur});
            }
      }
      vector<int>ans(m,0);
      sort(all.begin(),all.end());
      vector<bool>used(all.size(),0);
      for(int i=0;i<all.size();i++){
            if(used[i]) continue;
            int mx=-1;
            for(int j=i;j<all.size();j++){
                  if(all[i][0]!=all[j][0]) break;
                  used[j]=1;
                  mx=max(mx,j);
            }
            for(int j=i;j<=mx;j++){
                  int ii=all[j][1];
                  int node=all[j][2];
                  inf[node]=0;
            }
            for(int j=i;j<=mx;j++){
                  int ii=all[j][1];
                  int node=all[j][2];
                  if(I[ii]) inf[node]=1;
            }
            for(int j=i;j<=mx;j++){
                  int ii=all[j][1];
                  int node=all[j][2];
                  if(inf[node]) I[ii]=1;
                  ans[ii]+=I[ii];
            }
      }
      return ans;
}
