#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];
map<ll,vector<ll>>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);
                        }
                  }
            }
      }
      ll m=(int)S.size();
      vector<int>ans(m,0);
      set<ll>st;
      for(int i=0;i<m;i++) st.insert(D[i]),mp[D[i]].push_back(i);
      for(int i=0;i<m;i++) cc[i]=S[i];
      while(!st.empty()){
            ll mn=*st.begin();
            unordered_map<ll,ll>infc;
            for(int x:mp[mn]){
                  if(I[x]==1) infc[cc[x]]=1;
            }
            for(ll x:mp[mn]) if(infc[cc[x]]) I[x]=1;
            for(ll x:mp[mn]){
                  ans[x]+=I[x];
                  if(cc[x]==T[x]) continue;
                  cc[x]=go(cc[x],T[x]);
                  mp[mn+1].push_back(x);
                  st.insert(mn+1);
            }
            mp[mn].clear();
            st.erase(mn);
      }
      return {ans};
}
