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

const int N=5005;

struct dat {
    ll d;
    int s, t, inf, id;
} line[N];
bool cmp(dat &x, dat &y) {
    return x.inf>y.inf;
}

int _n, _m;
vector<int> g[N];
int par[20][N], dep[N], l[N], dist[N];

void dfs(int u, int p, int d) {
    dep[u]=d;
    par[0][u]=p;
    for(int v : g[u]) {
        if(v==p)
            continue;
        dfs(v, u, d+1);
    }
}
void init() {
    for(int i=1; i<20; i++) {
        for(int q=0; q<_n; q++) {
            if(par[i-1][q]==-1)
                par[i][q]=-1;
            else
                par[i][q]=par[i-1][par[i-1][q]];
        }
    }
    /**for(int i=0; i<20; i++) {
        for(int q=0; q<_n; q++) {
            cout << "for " << i << " " << q << " " << par[i][q] << endl;
        }
    }**/
}
int lca(int u, int v) {
    if(dep[u]>dep[v])
        swap(u, v);
    int diff=dep[v]-dep[u];
    ///cout << "lca qry: " << u << " " << v << endl;
    for(int i=19; i>=0; i--) {
        if(diff&(1<<i))
            diff-=(1<<i), v=par[i][v];
    }
    ///cout << "after equalizing: " << u << " " << v << endl;
    if(u==v) {
        ///cout << "result: " << u << endl;
        return u;
    }
    for(int i=19; i>=0; i--) {
        if(par[i][u]!=par[i][v])
            u=par[i][u], v=par[i][v];
    }
    ///cout << "result: " << par[0][u] << endl;
    return par[0][u];
}
bool f[N][N];
vector<int> paths[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
) {
    _n=R.size()+1;
    for(auto [u, v] : R)
        g[u].push_back(v), g[v].push_back(u);
    _m=D.size();
    vector<int> ans;
    for(int i=0; i<_m; i++)
        line[i]={D[i], S[i], T[i], I[i], i}, ans.push_back(0);
    for(int i=0; i<_n; i++)
        l[i]=2*_n;
    dfs(0, -1, 1);
    ///d=0;
    init();
    sort(line, line+_m, cmp);
    int last1=-1;
    for(int i=0; i<_m; i++) {
        int same=lca(line[i].s, line[i].t);
        int curr=line[i].s;
        dist[i]=dep[line[i].s]+dep[line[i].t]-2*dep[same]+1;
        if(line[i].inf==1) {
            last1=i;
        }
        while(curr!=same) {
            paths[i].push_back(curr);
            curr=par[0][curr];
        }
        int nw=line[i].t;
        vector<int> add;
        while(nw!=same) {
            add.push_back(nw);
            nw=par[0][nw];
        }
        add.push_back(same);
        reverse(add.begin(), add.end());
        for(int q : add)
            paths[i].push_back(q);
    }
    /**for(int i=0; i<_m; i++) {
        for(int q : paths[i])
            cout << q << " ";
        cout << endl;
    }**/
    for(int i=0; i<=last1; i++) {
        for(int q=0; q<paths[i].size(); q++)
            f[paths[i][q]][q]=1;
    }
    for(int t=0; t<_n; t++) {
        for(int i=0; i<_m; i++) {
            if(ans[line[i].id]!=0)
                continue;
            if(t<paths[i].size()) {
                if(f[paths[i][t]][t]) {
                    ans[line[i].id]=paths[i].size()-t;
                    for(int q=t; q<paths[i].size(); q++)
                        f[paths[i][q]][q]=1;
                }
            }
        }
    }
    return ans;
}
