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

const int MAXN = 5001;

vector<int> wut[MAXN];
vector<pair<int,int>> idk[MAXN];
vector<int> haha[MAXN];

int dude(int x, vector<pair<int,int>>& bruh) {
    if(bruh.empty()) {
        return -1;
    }
    int l = 0,r = bruh.size()-1;
    while(l < r) {
        int mid = (l+r)/2;
        if(bruh[mid].first < x) {
            l = mid+1;
        }
        else {
            r = mid;
        }
    }
    if(bruh[l].first == x) {
        return l;
    }
    else {
        return -1;
    }
}

bool dfs(int a, int t, int s, int k) {
    wut[s].push_back(a);
    if(a == k) {
        return true;
    }
    for(int v: haha[a]) {
        if(v != t) {
            if(dfs(v,a,s,k)) {
                return true;
            }
        }
    }
    wut[s].pop_back();
    return false;
}

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> z
) {
    int n = r.size()+1;
    int m = s.size();
    for(int i = 0; i < n-1; i++) {
        haha[r[i].first].push_back(r[i].second);
        haha[r[i].second].push_back(r[i].first);
    }
    for(int i = 0; i < m; i++) {
        dfs(s[i],-1,i,t[i]);
        for(int j = 0; j < wut[i].size(); j++) {
            idk[wut[i][j]].push_back({d[i]+j,i});
        }
    }
    priority_queue<pair<int,int>> banana;
    vector<int> ans(m,-1);
    for(int i = 0; i < m; i++) {
        if(z[i]) {
            banana.push({-d[i],i});
        }
    }
    while(!banana.empty()) {
        int x = banana.top().second;
        int br = -banana.top().first;
        banana.pop();
        if(ans[x] != -1) {
            continue;
        }
        ans[x] = br;
        for(int j = 0; j < wut[x].size(); j++) {
            if(j+d[x] < br) {
                continue;
            }
            int u = wut[x][j];
            int p = dude(j+d[x],idk[u]);
            if(p == -1) {
                continue;
            }
            if(idk[u][p].first == -1) {
                continue;
            }
            for(int k = p; k < idk[u].size(); k++) {
                if(idk[u][k].first != j+d[x]) {
                    break;
                }
                banana.push({-(j+d[x]),idk[u][k].second});
                idk[u][k] = {-1,-1};
            }
        }
    }
    for(int i = 0; i < m; i++) {
        if(ans[i] == -1) {
            ans[i] = 0;
        }
        else {
            ans[i] = wut[i].size()-(ans[i]-d[i]);
        }
    }
    return ans;
}
