#include "infection.h"
#include <vector>
#include <utility>

#include <bits/stdc++.h>

using namespace std;

//#include "grader.cpp"

#define nl '\n'

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;

    vector<vector<int>> adj(n);

    for(int i = 0; i < n - 1; i++) {
        int u = R[i].first, v = R[i].second;

        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    vector<int> par(n), tin(n), tout(n);

    int timer = 0;

    auto dfs = [&](auto &&self, int v, int p) -> void {
        par[v] = p;

        tin[v] = ++timer;

        for(auto to : adj[v]) {
            if(to == p) continue;

            self(self, to, v);
        }

        tout[v] = timer;
    };

    dfs(dfs, 0, 0);

    auto upper = [&](int u, int v) {
        return tin[v] >= tin[u] && tout[v] <= tout[u];
    };

    vector<vector<int>> where(n, vector<int>(n));

    for(int i = 0; i < n; i++) {
        int u = par[i], prev = i;

        while(prev != 0) {
            where[u][i] = prev;
            prev = u;
            u = par[u];
        }
    }

    int m = D.size();

    set<int> cur;

    vector<set<int>> mp(n);

    for(int i = 0; i < m; i++) {
        cur.emplace(i);
        mp[S[i]].emplace(i);
    }

    auto next = [&](int i) {
        if(S[i] == T[i]) return -1;
        if(!upper(S[i], T[i])) return par[S[i]];
        return where[S[i]][T[i]];
    };

    vector<int> ans(m);

    while(!cur.empty()) {
        vector<int> del;

//        cout << nl;

//        cout << "here: " << nl;

        for(auto i : cur) {
            if(!I[i]) {
                for(auto j : mp[S[i]]) {
                    if(I[j]) {
                        I[i] = true;
                        break;
                    }
                }
            }

//            cout << i << ' ';

            if(I[i]) ans[i]++;
        }

//        cout << nl;

//        for(auto i : cur) cout << S[i] << ' ';
//        cout << nl;

        for(auto i : cur) {
            int nxt = next(i);

            mp[S[i]].erase(i);

            if(nxt == -1) {
                del.push_back(i);
            } else {
                S[i] = nxt;
                mp[S[i]].emplace(i);
            }
        }

        for(auto i : del) cur.erase(i);
    }

    return ans;
}

/**

9 5
0 1
1 2
2 3
3 4
4 5
3 6
6 7
6 8
1 0 5 1
3 5 7 0
4 5 7 0
8 7 8 0
10 0 1 0


**/
