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

typedef long long ll;
typedef pair<ll, ll> pii;

constexpr int MAX = 150'005;

vector<int> G[MAX];
int col[MAX], dpth[MAX], jmp[MAX], par[MAX];

void dfs1(int v, int p, int d = 0) {
    dpth[v] = d;
    if (d) {
        par[v] = p;
        if (2 * dpth[jmp[p]] - dpth[p] == dpth[jmp[jmp[p]]])
            jmp[v] = jmp[jmp[p]];
        else
            jmp[v] = p;
    }
    else
        par[v] = jmp[v] = v;
    for (int u : G[v]) {
        if (u == p)
            continue;
        dfs1(u, v, d + 1);
    }
}

vector<tuple<ll, int, int>> E;

struct vertex {
    unordered_map<ll, unordered_set<int>> mp;
    size_t siz;
};

vector<pii> add1[MAX], add2[MAX], rem1[MAX], rem2[MAX];

pair<vertex, vertex> dfs2(int v, int p) {
    vertex fr, rev;
    auto insert1 = [&](ll t, int i) {
        if (!fr.mp[t].empty())
            E.emplace_back(t - dpth[v], i, *fr.mp[t].begin());
        // if (!rev.mp[t - 2 * dpth[v]].empty())
        //     E.emplace_back(t - dpth[v], i, *rev.mp[t - 2 * dpth[v]].begin());
        fr.mp[t].insert(i);
        fr.siz++;
    };
    auto insert2 = [&](ll t, int i) {
        if (!rev.mp[t].empty())
            E.emplace_back(t + dpth[v], i, *rev.mp[t].begin());
        // if (!fr.mp[t + 2 * dpth[v]].empty())
        //     E.emplace_back(t + dpth[v], i, *fr.mp[t + 2 * dpth[v]].begin());
        rev.mp[t].insert(i);
        rev.siz++;
    };
    for (auto [i, t] : add1[v])
        insert1(t, i);
    for (auto [i, t] : add2[v])
        insert2(t, i);
    for (int u : G[v]) {
        if (u == p)
            continue;
        auto [mp1, mp2] = dfs2(u, v);
        if (mp1.siz > fr.siz)
            swap(mp1, fr);
        for (auto [t, s] : mp1.mp) {
            for (int i : s)
                insert1(t, i);
        }
        if (mp2.siz > rev.siz)
            swap(mp2, rev);
        for (auto [t, s] : mp2.mp) {
            for (int i : s)
                insert2(t, i);
        }
    }
    for (auto [t, s] : fr.mp) {
        for (int i : s)
            if (!rev.mp[t - 2 * dpth[v]].empty())
                E.emplace_back(t - dpth[v], i, *rev.mp[t - 2 * dpth[v]].begin());
    }
    for (auto [i, t] : rem1[v])
        fr.mp[t].erase(i);
    for (auto [i, t] : rem2[v])
        rev.mp[t].erase(i);
    return {fr, rev};
}

vector<int> solve(
    vector<pair<int, int>> R,
    vector<ll> D,
    vector<int> S,
    vector<int> T,
    vector<bool> I
) {
    int N = R.size() + 1;
    for (auto [u, v] : R)
        G[u].push_back(v), G[v].push_back(u);
    dfs1(0, -1);
    int M = D.size();
    unordered_map<ll, vector<pii>> fr, rev;
    vector<ll> dist(M);
    for (int i = 0; i < M; i++) {
        int u = S[i], v = T[i];
        if (dpth[u] > dpth[v])
            swap(u, v);
        while (dpth[v] > dpth[u]) {
            if (dpth[jmp[v]] >= dpth[u])
                v = jmp[v];
            else
                v = par[v];
        }
        while (u != v) {
            if (jmp[u] != jmp[v])
                u = jmp[u], v = jmp[v];
            else
                u = par[u], v = par[v];
        }
        dist[i] = dpth[S[i]] + dpth[T[i]] - 2 * dpth[v];
        add1[S[i]].push_back({i, D[i] + dpth[S[i]]});
        add2[T[i]].push_back({i, D[i] + dist[i] - dpth[T[i]]});
        rem1[v].push_back({i, D[i] + dpth[S[i]]});
        rem2[v].push_back({i, D[i] + dist[i] - dpth[T[i]]});
    }
    dfs2(0, -1);
    sort(E.begin(), E.end());
    vector<int> ans(M);
    priority_queue<pii> pq;
    for (int i = 0; i < M; i++) {
        if (I[i])
            ans[i] = dist[i] + 1, pq.push({-D[i], i});
    }
    vector<vector<pii>> G2(M);
    for (auto [w, u, v] : E) {
        G2[u].push_back({v, w});
        G2[v].push_back({u, w});
    }
    vector<bool> vis(M);
    while (!pq.empty()) {
        auto [d, v] = pq.top();
        d *= -1;
        pq.pop();
        if (vis[v])
            continue;
        vis[v] = true;
        if (!I[v]) {
            I[v] = true;
            ans[v] = dist[v] + 1 - (d - D[v]);
        }
        for (auto [u, w] : G2[v]) {
            if (!vis[u] && w >= d)
                pq.push({-w, u});
        }
    }
    return ans;
}
