#include "infection.h"
#include <bits/extc++.h>
using namespace std;
using namespace __gnu_pbds;

#define int long long

const int N = 5005;
vector<int> g[N];
int par[N];

void dfs(int u, int p = 0) {
    par[u] = p;
    for (int v : g[u]) {
        if (v == p) {
            continue;
        }
        dfs(v, u);
    }
}

struct event {
    int time, node, vip;

    bool operator<(event other) {
        if (time != other.time) {
            return time < other.time;
        }
        return node < other.node;
    }
};

#undef int

vector<int> solve(vector<pair<int, int>> R, vector<long long> D, vector<int> S, vector<int> T, vector<bool> I) {
    int n = R.size() + 1;
    for (auto [u, v] : R) {
        u++, v++;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    int m = D.size();
    vector<int> ans(m);
    vector<event> ev;
    #define int long long
    vector<int> Path[m];
    for (int i = 0; i < m; i++) {
        S[i]++, T[i]++;
        dfs(S[i]);
        vector<int> path;
        int u = T[i];
        while (u) {
            path.push_back(u);
            u = par[u];
        }
        reverse(path.begin(), path.end());
        for (int j = 0; j < path.size(); j++) {
            ev.push_back({D[i] + j, path[j], i});
        }
        Path[i] = path;
    }
    sort(ev.begin(), ev.end());
    gp_hash_table<int, bool> inf[n + 1];
    for (int id = 0; id < ev.size(); id++) {
        auto [time, node, vip] = ev[id];
        if (time) {
            auto it = lower_bound(Path[vip].begin(), Path[vip].end(), node);
            if (it != Path[vip].begin() and time) {
                if (inf[*--it][time - 1]) {
                    I[vip] = true;
                }
            }
        }
        if (I[vip]) {
            inf[node][time] = true;
        }
    }
    for (auto [time, node, vip] : ev) {
        ans[vip] += inf[node][time];
    }
    return ans;
}
