#include "infection.h"
#include <algorithm>
#include <cstdio>
using namespace std;

#define MAXN 150000
#define MAXM 150000

vector<int> adj[MAXN];
int depth[MAXN],parent[MAXN];
vector<int> path[MAXM];
vector<pair<int,int>> ids[MAXN],all;
long long infected[MAXM];

void dfs(int u,int p,int d){
  depth[u]=d;
  parent[u]=p;
  for(int v:adj[u]){
    if(v!=p){
      dfs(v,u,d+1);
    }
  }
}

vector<int> getPath(int st,int dr){
  vector<int> leftp,rightp;

  while(depth[st]<depth[dr]){
    rightp.push_back(dr);
    dr=parent[dr];
  }

  while(depth[dr]<depth[st]){
    leftp.push_back(st);
    st=parent[st];
  }

  while(st!=dr){
    leftp.push_back(st);
    st=parent[st];
    rightp.push_back(dr);
    dr=parent[dr];
  }
  leftp.push_back(st);
  for(auto it=rightp.rbegin();it!=rightp.rend();it++){
    leftp.push_back(*it);
  }

  return leftp;
}

vector<int> solve(vector<pair<int, int>> R,vector<long long> D,vector<int> S,vector<int> T,vector<bool> I){
  int m=D.size(),i,n=R.size()+1;
  for(auto p:R){
    adj[p.first].push_back(p.second);
    adj[p.second].push_back(p.first);
  }
  dfs(0,-1,0);

  for(i=0;i<m;i++){
    infected[i]=-1;
    if(I[i]){
      infected[i]=D[i];
    }
    all.push_back({D[i],i});
    path[i]=getPath(S[i],T[i]);
    long long t=D[i];
    for(int u:path[i]){
      ids[u].push_back({t,i});
      t++;
    }
  }
  sort(all.begin(),all.end());

  for(int u=0;u<n;u++){
    sort(ids[u].begin(),ids[u].end());
  }

  vector<int> ans(m);
  for(auto p:all){
    i=p.second;
    long long len=0;
    if(I[i]){
      len=path[i].size()-(infected[i]-D[i]);
      long long tin=D[i];
      for(int u:path[i]){
        auto it=ids[u].begin();
        // fprintf(stderr,"%d intra in %d la timpul %lld\n",i,u,tin);
        while(it!=ids[u].end()&&it->first!=tin){
          it++;
        }

        while(it!=ids[u].end()&&it->first==tin&&it->second<=i){
          it++;
        }

        while(it!=ids[u].end()&&it->first==tin){
          int id=it->second;
          I[id]=true;
          infected[id]=tin;
          // fprintf(stderr,"%d la infectat si pe %d la momentul %lld si orasul %d\n",i,id,tin,u);
          it++;
        }
        tin++;
      }
    }
    ans[i]=len;
  }

  return ans;
}
