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

using ll = long long;

vector<int> solve(vector<pair<int, int>> edges, vector<long long> day, vector<int> s, vector<int> t, vector<bool> I) {
    int n = (int)edges.size() + 1, m = (int)day.size();
    vector<int> ans(m);
    vector<vector<int>> g(n + 1);
    for(int i = 0; i < n - 1; i++) {
        auto [x, y] = edges[i];
        g[x].push_back(y);
        g[y].push_back(x);
    }
    vector<int> p(n);
    function<void(int)> dfs = [&](int v) {
        for(auto u : g[v]) {
            if(u == p[v]) continue;
            p[u] = v;
            dfs(u);
        }
    };
    vector<array<ll, 3>> qw;
    for(int i = 0; i < m; i++) {
        p[s[i]] = s[i];
        dfs(s[i]);
        vector<int> v;
        int st = t[i];
        while(true) {
            v.push_back(st);
            if(st == s[i]) break;
            st = p[st];
        }
        reverse(v.begin(), v.end());
        for(int j = 0; j < (int)v.size(); j++) {
            qw.push_back({day[i] + j, v[j], i});
        }
    }
    sort(qw.begin(), qw.end());
    for(int i = 0; i < (int)qw.size(); i++) {
        int LST = i;
        ll tm = qw[i][0];
        while(i < (int)qw.size() && qw[i][0] == tm) {
            i++;
        }
        i--;
        for(int j = LST; j <= i; j++) {
            tm = qw[j][1];
            int lst = j;
            bool flag = 0;
            while(j <= i && tm == qw[j][1]) {
                if(I[qw[j][2]]) flag = 1;
                j++;
            }
            j--;
            if(flag) {
                for(int k = lst; k <= j; k++) {
                    I[qw[k][2]] = 1;
                    ans[qw[k][2]]++;
                }
            }
        }

    }
    return ans;

};
