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

#define ll long long

const int MN = 5005;
const int MM = 5005;

int n, m;

vector<vector<int>> tree(MN), curr(MN);
int par[MN];
bool sub[MN][MN]; 
int ans[MM], dest[MM];
bool inf[MM], ok[MM];

void prec(int node, int p) {
    par[node] = p;
    sub[node][node] = true;
    for (auto next: tree[node]) {
        if (next == p) continue;
        prec(next, node);
        for (int i = 0; i < n; i++) {
            sub[node][i] = sub[node][i] || sub[next][i];
        }
    }
}

void infect() {
    for (int i = 0; i < n; i++) {
        bool bad = false;
        for (auto c: curr[i]) {
            if (ok[c]) continue;
            if (inf[c]) bad = true;
        }
        if (bad) {
            for (auto c: curr[i]) {
                if (ok[c]) continue;
                inf[c] = true;
            }
        }
    }
}

void add_ans() {
    infect();
    for (int i = 0; i < m; i++) {
        if (ok[i]) continue;
        if (inf[i]) ans[i]++;
    }
}

int move() {
    int tot = 0;
    vector<vector<int>> nxt(n);
    for (int i = 0; i < n; i++) {
        for (auto c: curr[i]) {
            if (ok[c]) continue;
            tot++;
            if (dest[c] == i) {
                ok[c] = true;
                continue;
            }
            if (!sub[i][dest[c]]) {
                nxt[par[i]].push_back(c);
            }
            else {
                for (auto nn: tree[i]) {
                    if (nn == par[i]) continue;
                    if (sub[nn][dest[c]]) {
                        nxt[nn].push_back(c);
                    }
                }
            }
        }
    }
    curr = nxt;
    return tot;
}

vector<int> solve(vector<pair<int, int>> R, vector<ll> D, vector<int> S, vector<int> T, vector<bool> I) {
    n = R.size() + 1;
    m = I.size();
    tree.resize(n+1);
    for (auto [u, v]: R) {
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    prec(0, -1);
    for (int i = 0; i < m; i++) {
        inf[i] = I[i];
    }
    vector<pair<ll, int>> allD;
    for (int i = 0; i < m; i++) {
        allD.push_back({D[i], i});
        dest[i] = T[i];
    }
    sort(allD.rbegin(), allD.rend());
    int currD = 0;
    while (!allD.empty()) {
        while (!allD.empty() && allD.back().first <= currD) {
            int id = allD.back().second;
            allD.pop_back();
            curr[S[id]].push_back(id);
        }
        add_ans();
        int ret = move();
        if (ret) {
            currD++;
            continue;
        }
        currD = allD.back().first;
    }
    for (int i = 0; i < n; i++) {
        add_ans();
        move();
    }
    vector<int> ret;
    for (int i = 0; i < m; i++) {
        ret.push_back(ans[i]);
    }
    return ret;
}
