#include "infection.h"
#include <vector>
#include <utility>
#include <set>
#include <iostream>
#include <cassert>
using namespace std;

const int N = 200000 + 10;

vector<int> adj[N];

int ans[N], depth[N], up[N][20];
bool marked[5005];
int mat[5005][5005];

void dfs0(int node, int parent) {
    up[node][0] = parent;
    for (int i = 1; i < 20; i++) {
        up[node][i] = up[up[node][i - 1]][i - 1];
    }
    for (auto i : adj[node]) {
        if (i == parent) continue;
        depth[i] = depth[node] + 1;
        dfs0(i, node);
    }
}

int lca(int u, int v) {
    if (depth[u] < depth[v]) swap(u, v);
    int d = depth[u] - depth[v];
    for (int i = 0; i < 20; i++) {
        if (d & (1 << i)) {
            u = up[u][i];
        }
    }
    if (u == v) return u;
    for (int i = 19; i >= 0; i--) {
        if (up[u][i] != up[v][i]) {
            u = up[u][i];
            v = up[v][i];
        }
    }
    return up[u][0];
}

void fill(int node, int parent, int& t, bool fl, int i) {
    if (!fl) {
        while (node != parent) {
            mat[i][t] = node;
            node = up[node][0];
            t++;
        }
        mat[i][t] = node;
        return;
    }
    while (node != parent) {
        mat[i][t] = node;
        t--;
        node = up[node][0];
    }
}

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) {
    for (auto e : R) {
        int u = e.first;
        int v = e.second;
        u++, v++;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs0(1, 1);
    vector<int> res;
    for (int i = 0; i < D.size(); i++) {
        for (int j = 0; j < R.size() + 2; j++) {
            mat[i][j] = -1;
        }
    }
    for (int i = 0; i < D.size(); i++) {
        if (I[i]) marked[i] = 1;
        int lc = lca(S[i] + 1, T[i] + 1);
        int t = 0, d = depth[S[i] + 1] + depth[T[i] + 1] - 2*depth[lc];
        fill(S[i] + 1, lc, t, 0, i);
        fill(T[i] + 1, lc, d, 1, i);
    }
    for (int i = 0; i < R.size() + 2; i++) {
        bool fl = 0;
        set<int> var;
        for (int j = 0; j < D.size(); j++) {
            if (mat[j][i] == -1) continue;
            if (marked[j]) {
                var.insert(mat[j][i]);
            }
        }
        for (int j = 0; j < D.size(); j++) {
            if (mat[j][i] == -1) continue;
            if (var.count(mat[j][i])) {
                marked[j] = 1;
                ans[j]++;
            }
        }
    }
    for (int i = 0; i < D.size(); i++) {
        res.push_back(ans[i]);
    }
    return res;
}
