#include "infection.h"
#include <bits/stdc++.h>
#define pii pair<int,int>
#define ff first
#define ss second
using namespace std;
using ll = long long;
ll inf = 1e18;
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> bad
) {
    int q = s.size();
    int n = r.size()+1;
    vector<ll> ds = d;
    sort(ds.begin(),ds.end());

    ds.push_back(inf);
    vector<int> res(q,0);
    vector<vector<int>> p(n,vector<int>(n));
    vector<vector<int>> g(n);
    for(pii a : r)
        g[a.ff].push_back(a.ss),
        g[a.ss].push_back(a.ff);
    for(int rt = 0;rt < n;rt++)
    {
        function<void(int,int)> dfs = [&](int cur,int prev)
        {
            p[rt][cur] = prev;
            for(int a : g[cur])
                if(a!=prev)
                    dfs(a,cur);
        };
        dfs(rt,rt);
        p[rt][rt] = -1;
    }
    vector<int> loc(q,-1);
    map<ll,vector<int>> poc;
    for(int i = 0; i < q; i++)
        poc[d[i]].push_back(i);
    ll dan = -1;
    int act = 0;
    while(dan<inf)
    {
        
        if(act==0)
            dan = *upper_bound(ds.begin(),ds.end(),dan);
        else
            dan = dan+1;
        for(int a:poc[dan])
            loc[a] = s[a],act++;
        vector<int> bol(n,0);
        for(int i = 0; i < q; i++)
            if(loc[i]!=-1)
                bol[loc[i]]|=int(bad[i]);
        for(int i = 0; i < q; i++)
            if(loc[i]!=-1)
                bad[i]=bol[loc[i]];
        for(int i = 0; i < q; i++)
            if(loc[i]!=-1)
                res[i]+=bad[i];
        for(int i = 0; i < q; i++)
            if(loc[i]!=-1)
            {
                loc[i] = p[t[i]][loc[i]];
                act+=(loc[i]==-1);
            }
            
    }
    return res;
}