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

vector <vector <int> > gr;
vector <vector <int> > binlift;
vector <int> in;
vector <int> out;
vector <int> parent;
int timer = 0;

using llint = long long;

void dfs (int p, int par) {
    parent[p] = par;
    in[p] = out[p] = ++timer;

    binlift[p][0] = par;
    for (int i = 1; i < 20; i++)
        binlift[p][i] = binlift[binlift[p][i - 1]][i - 1];

    for (int& e : gr[p])
        if (e != par)
            dfs (e, p);
    out[p] = timer;
}

inline bool is_ancestor (int a, int b) {
    return in[a] <= in[b] && in[b] <= out[a];
}

int lca (int a, int b) {
    if (is_ancestor (a, b))
        return a;
    if (is_ancestor (b, a))
        return b;
    for (int i = 19; i >= 0; i--)
        if (!is_ancestor (binlift[a][i], b))
            a = binlift[a][i];
    return binlift[a][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
) {
    int N = R.size() + 1;
    gr.resize (N + 1);
    binlift.resize (N + 1, vector <int> (20));
    in.resize (N + 1);
    out.resize (N + 1);
    parent.resize (N + 1);
    out[0] = 2e9;
    for (auto [a, b] : R) {
        a++, b++;
        gr[a].emplace_back (b);
        gr[b].emplace_back (a);
    }

    dfs (1, 0);

    for (int& e : S)
        e++;
    for (int& e : T)
        e++;

    vector <vector <int> > routes (S.size());
    for (int i = 0; i < S.size(); i++) {
        int a = S[i], b = T[i];
        int c = lca (a, b);
        vector <int>& route = routes[i];

        while (a != c) {
            route.emplace_back (a);
            a = parent[a];
        }
        route.emplace_back (c);
        
        vector <int> temp;
        while (b != c) {
            temp.emplace_back (b);
            b = parent[b];
        }
        reverse (temp.begin(), temp.end());
        for (int& e : temp)
            route.emplace_back (e);
    }

    llint t = -1e18;
    set <int> todo;
    vector <int> curind (S.size());
    multiset <pair <llint, int> > s;
    vector <int> ans (S.size());
    for (int i = 0; i < S.size(); i++)
        s.emplace (D[i], i);
    while (!s.empty()) {
        while (s.begin()->first == t) {
            todo.emplace (s.begin()->second);
            s.erase (s.begin());
        }

        if (todo.size() == 0) {
            todo.emplace (s.begin()->second);
            t = s.begin()->first;
            s.erase (s.begin());
        }

        // cout << t << ": \n";
        // cout << "todo: ";
        // for (auto&& e : todo)
        //     cout << e << ' ';
        // cout << '\n';
    

        vector <vector <int> > ids (N + 1);
        vector <bool> ill_add (N + 1);
        for (auto&& e : todo) {
            ids[routes[e][curind[e]]].emplace_back (e);
            if (I[e])
                ill_add[routes[e][curind[e]]] = 1;
        }
        for (int i = 1; i <= N; i++)
            if (ill_add[i])
                for (auto&& e : ids[i])
                    I[e] = 1;
        for (auto&& e : todo)
            if (I[e])
                ans[e]++;
        vector <int> to_rem;
        for (auto&& e : todo) {
            curind[e]++;
            if (curind[e] >= routes[e].size())
                to_rem.emplace_back (e);
        }
        for (auto&& e : to_rem)
            todo.erase (e);
        t++;
    }
    
    return ans;
}
