#include "infection.h"
#include <vector>
#include <utility>
#include<iostream>
#include<queue>
#include<map>
using namespace std;
vector<vector<int>>dist;
vector<vector<int>>gp;
void dfs(int u, int p, int skizb) {

    for (int x : gp[u]) {
        if (x != p) {
            dist[skizb][x] = dist[skizb][u] + 1;
            dfs(x, u, skizb);
        }
    }


}
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;
    int m = (D.size());
    gp = vector<vector<int>>(n);
    dist = vector<vector<int>>(n, vector<int>(n));
    for (int i = 0; i < n - 1; i++) {
        gp[R[i].first].push_back(R[i].second);
        gp[R[i].second].push_back(R[i].first);

    }
    for (int i = 0; i < n; i++) {
        dfs(i, -1, i);
    }
    priority_queue<pair<pair<int, int>, pair<int, int>>>q;
    for (int i = 0; i < m; i++) {
        q.push({ { -D[i],S[i]},{I[i],i} });
    }
    int last = -1, f = -1, gag = -1;
    vector<int>ans(m);
    map<pair<int, int>, int>mp;
    while (q.size())
    {
        int d = -q.top().first.first, arj = q.top().second.first, u = q.top().first.second, cur = q.top().second.second;
        q.pop();
        //mp[{}]
        if (arj == 1) {
            ans[cur]++;
        }
        else {
            if (d == last && u == gag && f == 1) {
                arj = 1;
                I[cur] = 1;
                ans[cur]++;
            }
        }
        last = d;
        f = arj;
        gag = u;
        if (u == T[cur])continue;
        int ur = -1;
        for (int x : gp[u]) {
            if (dist[x][T[cur]] <= dist[u][T[cur]])ur = x;
        }
        q.push({ {-(d + 1),ur},{arj,cur} });
    }

    return ans;



}
