#include "infection.h"
#include <vector>
#include <utility>
#include <cstring>
#include <queue>
using namespace std;

vector<vector<int>> tree, paths;
vector<int> ans;
bool vis[150000], fp;
vector<long long> infs[150000];

void dfs(int s, int t, int i, int p = -1) {
    if (s == t || fp) {
        fp = true;
        return;
    }
    vis[s] = true;
    for (const auto nxt : tree[s]) {
        if (vis[nxt] || fp)
            continue;
        paths[i].push_back(nxt);
        dfs(nxt, t, i, s);
        if (fp)
            break;
        paths[i].pop_back();
    }
}

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, m = D.size();
    tree.resize(n), paths.resize(m), ans.resize(m);
    for (int i = 0; i < n - 1; i++)
        tree[R[i].first].push_back(R[i].second), tree[R[i].second].push_back(R[i].first);
    for (int i = 0; i < m; i++) {
        paths[i].push_back(S[i]);
        memset(vis, false, sizeof(vis)), fp = false;
        dfs(S[i], T[i], i);
    }
    queue<int> iis;
    for (int i = 0; i < m; i++)
        if (I[i])
            iis.push(i);
    bool ch;
    do {
        ch = false;
        while (!iis.empty()) {
            int i = iis.front();
            iis.pop();
            ans[i] = paths[i].size();
            for (int j = 0; j < paths[i].size(); j++)
                infs[paths[i][j]].push_back(D[i] + j);
        }
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < paths[i].size(); j++) {
                for (int p = 0; p < infs[paths[i][j]].size(); p++) {
                    if (infs[paths[i][j]][p] == D[i] + j && ans[i] < int(paths[i].size() - j - 1))
                        ans[i] = int(paths[i].size() - j - 1), ch = true, iis.push(i);
                }
            }
        }
    } while (ch);
    return ans;
}
