#include "infection.h"
#include <bits/stdc++.h>
#define NMAX 5000
#define LOG 20
using namespace std;

vector<int> tree[NMAX];
int n, m, depth[NMAX], anc[LOG][NMAX];

void dfs(int node, int parent)
{
    anc[0][node]=parent;
    for(auto& next : tree[node])
    {
        if(next==parent) continue;
        depth[next]=1+depth[node];
        dfs(next, node);
    }
}

void precalc()
{
    for(int e=1; e<LOG; e++)
    {
        for(int i=0; i<n; i++)
            anc[e][i]=anc[e-1][anc[e-1][i]];
    }
}

int up(int node, int k)
{
    for(int e=LOG-1; e>=0; e--)
    {
        if((1<<e)<=k)
        {
            k-=(1<<e);
            node=anc[e][node];
        }
    }
    return node;
}

int lca(int node1, int node2)
{
    if(depth[node1]<depth[node2]) swap(node1, node2);
    node1=up(node1, depth[node1]-depth[node2]);
    if(node1==node2) return node1;
    for(int e=LOG-1; e>=0; e--)
    {
        if(anc[e][node1]!=anc[e][node2])
        {
            node1=anc[e][node1];
            node2=anc[e][node2];
        }
    }
    return anc[0][node1];
}

struct data
{
    int day=0, node=0, idx=0, is_inf=0;
};

bool cmp(data a, data b)
{
    if(a.day!=b.day) return a.day<b.day;
    if(a.node!=b.node) return a.node<b.node;
    return a.is_inf;
}

vector<int> solve(vector<pair<int, int>> r, vector<long long> d, vector<int> s, vector<int> t, vector<bool> inf)
{
    n=r.size()+1;
    m=d.size();
    for(auto& [u, v] : r)
    {
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    dfs(0, -1);
    precalc();
    vector<data> op;
    for(int i=0; i<m; i++)
    {
        vector<int> path, path_t;
        int _lca=lca(t[i], s[i]);
        for(int node=t[i]; node!=_lca; node=anc[0][node]) path_t.push_back(node);
        for(int node=s[i]; node!=_lca; node=anc[0][node]) path.push_back(node);
        path.push_back(_lca);
        while(!path_t.empty())
        {
            path.push_back(path_t.back());
            path_t.pop_back();
        }
        for(int j=0; j<path.size(); j++)
        {
            op.push_back({d[i]+j, path[j], i, inf[i]});
        }
    }
    sort(op.begin(), op.end(), cmp);
    vector<int> ans(m, 0);
    bool infected=0;
    for(int i=0; i<op.size(); i++)
    {
        op[i].is_inf=inf[op[i].idx];
        if(op[i].is_inf) infected=1;
        inf[op[i].idx]=infected;
        ans[op[i].idx]+=inf[op[i].idx];
        if(i==op.size()-1 || (op[i].day!=op[i+1].day || op[i].node!=op[i+1].node)) infected=0;
    }
    return ans;
}