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

const long long MAXN = 5001;

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

long long dude(long long x, vector<pair<long long,long long>>& bruh) {
    if(bruh.empty()) {
        return -1;
    }
    long long l = 0,r = bruh.size()-1;
    while(l < r) {
        long long 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(long long a, long long t, long long s, long long k) {
    wut[s].push_back(a);
    if(a == k) {
        return true;
    }
    for(long long 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
) {
    long long n = r.size()+1;
    long long m = s.size();
    for(long long i = 0; i < n-1; i++) {
        haha[r[i].first].push_back(r[i].second);
        haha[r[i].second].push_back(r[i].first);
    }
    vector<long long> apple(m,LLONG_MAX);
    for(long long i = 0; i < m; i++) {
        dfs(s[i],-1,i,t[i]);
        for(long long j = 0; j < wut[i].size(); j++) {
            idk[wut[i][j]].push_back({d[i]+j,i});
        }
    }
    for(long long i = 0; i < n; i++) {
        sort(idk[i].begin(),idk[i].end());
    }
    vector<long long> ans(m,LLONG_MAX);

    for(long long i = 0; i < m; i++) {
        if(z[i]) {
            apple[i] = d[i];
        }
    }
    while(true) {
        pair<long long,long long> sm = {LLONG_MAX,-1};
        for(long long i = 0; i < m; i++) {
            if(apple[i] != -1) {
                sm = min(sm,{apple[i],i});
            }
        }
        if(sm.first == LLONG_MAX) {
            break;
        }
        long long x = sm.second;
        long long br = sm.first;
        ans[x] = br;
        apple[x] = -1;
        for(long long j = 0; j < wut[x].size(); j++) {
            if(j+d[x] < br) {
                continue;
            }
            long long u = wut[x][j];
            long long p = dude(j+d[x],idk[u]);
            if(p == -1) {
                continue;
            }
            if(idk[u][p].second == -1) {
                continue;
            }
            for(long long k = p; k < idk[u].size(); k++) {
                if(idk[u][k].second == -1 || idk[u][k].first != j+d[x]) {
                    break;
                }
                apple[idk[u][k].second] = min((long long)apple[idk[u][k].second],j+d[x]);
                idk[u][k] = {idk[u][k].first,-1};
            }
        }
    }
    for(long long i = 0; i < m; i++) {
        if(ans[i] == LLONG_MAX) {
            ans[i] = 0;
        }
        else {
            ans[i] = wut[i].size()-(ans[i]-d[i]);
        }
    }
    vector<int> ans2(m);
    for(int i = 0; i < m; i++) {
        ans2[i] = ans[i];
    }
    return ans2;
}
