#include "infection.h"
#include <bits/stdc++.h>


const int MAXN = 2e5 + 5;
const int KLOG = 18;
int n,m,dep[MAXN],p[MAXN],jump[MAXN][KLOG];
std::vector<int>g[MAXN];
void dfs(int beg)
{
    dep[beg] = dep[p[beg]] + 1;
    for(int nb : g[beg])
    {
        if(nb == p[beg]) continue;
        p[nb] = beg;
        jump[nb][0] = beg;
        dfs(nb);
    }
}

void pre_jump()
{
    for(int j = 1; j < KLOG; j++)
    {
        //cout << "forj" << j << " : ";
        for(int i = 0; i < n; i++)
        {
            jump[i][j] = jump[jump[i][j-1]][j-1];
           //cout << jump[i][j] << " ";
        }
        //cout << endl;
    }
}

int lca(int a, int b)
{
    int raz = abs(dep[a] - dep[b]);
    for(int j = 0; j < KLOG; j++)
    {
        if(raz&(1<<j))
        {
            if(dep[a] > dep[b]) a = jump[a][j];
            else b = jump[b][j];
        }
    }
    if(a == b) return a;
    for(int j = KLOG-1; j >= 0; j--)
    {
        if(jump[a][j] != jump[b][j])
        {
            a = jump[a][j];
            b = jump[b][j];
        }
    }
    return jump[a][0];
}
int dist(int a, int b)
{
    return dep[a] + dep[b] - 2*dep[lca(a,b)];
}


long long ver[MAXN],day[MAXN];
struct Person
{
    long long day;
    int ver,j;
};
Person sp[25000000];
int spsz;
bool cmpsp(Person p1,  Person p2)
{
    if(p1.day != p2.day) return p1.day < p2.day;
    if(p1.ver != p2.ver) return p1.ver < p2.ver;
    return p1.j < p2.j;
}

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)
{
    m = D.size();
    for(std::pair<int,int>x : R)
    {
        n = std::max(n,std::max(x.first,x.second));
        g[x.first].push_back(x.second);
        g[x.second].push_back(x.first);
    }
    n++;/// <n
    //cout << "N is : " << n << endl;

    p[0] = jump[0][0] = 0;
    dfs(0);
    pre_jump();

    ///sweep point
    std::vector<std::pair<long long, int> > sw;
    for(int i = 0; i < m; i++)
    {
        sw.push_back({D[i],0});
        sw.push_back({D[i] + dist(S[i],T[i]),1});
    }
    std::sort(sw.begin(),sw.end());
    std::unordered_map<long long,int>mymap;
    int midx = 0;
    long long st = sw[0].first;
    int bal = 1;
    for(int j = 1; j < sw.size(); j++)
    {
        if(sw[j].second == 0) bal++;
        else bal--;

        if(bal == 0)
        {
            for(long long i = st; i <= sw[j].first; i++) mymap[i] = midx++;
            if(j+1 < sw.size()) st = sw[j+1].first;///for next section
        }
    }


    for(int j = 0; j < m; j++) ver[j] = S[j];
    for(int j = 0; j < m; j++) day[j] = D[j];
    for(int q = 1; q <= n; q++)
    {
        for(int j = 0; j < m; j++)
        {
            if(ver[j] == -1) continue;

            sp[spsz++] = {mymap[day[j]],ver[j],j};

            ///next
            if(ver[j] == T[j]) ver[j] = -1;
            else
            {
                int ve = ver[j];
                for(int nb : g[ve]) if(dist(nb,T[j]) == dist(ve,T[j]) - 1) {ver[j] = nb; break;}
                day[j] ++;
            }
        }
    }
    std::sort(sp,sp+spsz,cmpsp);


    std::vector<int>ans,in;
    ans.resize(m);
    in.resize(m);
    for(int j = 0; j < m; j++) in[j] = I[j];

    //for(Person x : sp) cout << x.day << " " << x.ver << " " << x.j << " inaperson" << endl;
    for(int i = 0; i < spsz; i++)
    {
        int j = i;
        while(j+1 < spsz && sp[j+1].day == sp[j].day && sp[j+1].ver == sp[j].ver) j++;

        bool f = false;
        for(int k = i; k <= j; k++) if(in[sp[k].j]) f = true;
        if(f) for(int k = i; k <= j; k++) in[sp[k].j] = 1;
        for(int k = i; k <= j; k++) if(in[sp[k].j]) ans[sp[k].j]++;


        i = j;
    }
    return ans;
}

