#include "infection.h"
#include <vector>
#include <iostream>
#include <algorithm>
#include <utility>
#include <unordered_set>

using namespace std;
const int MAX_N = 150042;

int par[MAX_N][20], depth[MAX_N];
vector<int> adj_list[MAX_N];

void dfs(int u, int p, int d){
    depth[u] = d;
    par[u][0] = p;
    for(int v : adj_list[u])
        if(v != p) // tree
            dfs(v, u, d + 1);
}

int lift_k_up(int v, int k){
    for(int i = 0; i < 20; ++i)
        if(k & (1 << i))
            v = par[v][i];
    return v;
}

int LCA(int a, int b){
    if(depth[a] < depth[b]) swap(a, b); // we need a to be deeper than be
    a = lift_k_up(a, depth[a] - depth[b]); // a and b are on the same depth
    if(a == b) return a; // haha this whoopsie was missing
    for(int i = 19; i >= 0; --i){
        if(par[a][i] != par[b][i]){
            a = par[a][i];
            b = par[b][i];
        }
    }
    return par[a][0];
}

std::vector<int> solve(std::vector<std::pair<int, int>> R, std::vector<long long> D, std::vector<int> S, std::vector<int> T, std::vector<bool> I) {
    vector<int> ans;
    int n = R.size() + 1, m = D.size();
    for(int i = 0; i < n - 1; ++i){
        int u = R[i].first, v = R[i].second;
        adj_list[u].push_back(v);
        adj_list[v].push_back(u);
    }
    dfs(0, 0, 0);
    for(int i = 1; i < 20; ++i)
        for(int j = 0; j < n; ++j)
            par[j][i] = par[par[j][i - 1]][i - 1];
    vector<int> paths[m];
    // we calculate the path of each of these fuckers
    int longest_path = 0;
    for(int i = 0; i < m; ++i){
        int mid = LCA(S[i], T[i]);
        // their path is from S to LCA
        int v = S[i];
        while(v != mid){
            paths[i].push_back(v);
            v = par[v][0];
        }
        paths[i].push_back(mid);
        // then from the end to the LCA but in reverse
        int temp = paths[i].size();
        v = T[i];
        while(v != mid){
            paths[i].push_back(v);
            v = par[v][0];
        }
        reverse(paths[i].begin() + temp, paths[i].end());
        if(paths[i].size() + D[i] > longest_path) longest_path = paths[i].size() + D[i];
        ans.push_back(0);
        //cout << "Path of guy " << i << ": "; for(int p : paths[i]) cout << p << " "; cout << endl;
    }
    // for the cases where they all start traveling day 0
    // this should be running in something like O(N*M), yeah?
    vector<int> travel_day;
    for(int i = 0; i < m; ++i) travel_day.push_back(D[i]);
    for(int day = 0; day < longest_path; ++day){
       // cout << "DAY " << day << endl;
        unordered_set<int> infected_cities;
        for(int person = 0; person < m; ++person){
            if(day >= D[person] && paths[person].size() + D[person] <= day) continue; // they have left or are yet to arrive
            if(I[person] == 1){ // infected fucker
              //  cout <<"\t"<< person << " is from the infected and he is currently in city " << paths[person][day] << endl;
                ans[person]++;
                infected_cities.insert(paths[person][travel_day[person]]);
                travel_day[person]++;
            }
        }
        // now we know all the infected cities, we can check if there are other people that get infected
        for(int person = 0; person < m; ++person){
            if(day >= D[person] && paths[person].size() + D[person] <= day) continue; // they have left or are yet to arrive
            if(I[person] == 0){ // uninfected fucker
                if(infected_cities.contains(paths[person][travel_day[person]])){
                    ans[person]++;
                    I[person] = 1;
                    travel_day[person]++;
                }
            }
        }
    }
    return ans;
}
