#include "infection.h"
#include <algorithm>
#include <iostream>
#include <cassert>
#include <numeric>
#include <random>
#include <stack>
#include <queue>
#include <set>
#include <map>

const int MAXN = 5000 + 10;
const int MAXLOG = 20;
const int INTINF = 1e9 + 10;

int n, m, k;
std::vector<int> ans;

int pathLen[MAXN];
int isBad[MAXN][MAXN];
int position[MAXN][MAXN];
std::vector<int> g[MAXN];
int parent[MAXN][MAXLOG];
int depth[MAXN];
int in[MAXN];
int out[MAXN];
int timer;

bool isAnc(int u, int v)
{
    return in[u] <= in[v] && out[u] >= out[v];
}

int getLCA(int u, int v)
{
    if(isAnc(u, v))
    {
        return u;
    }

    if(isAnc(v, u))
    {
        return v;
    }

    for(int lg = MAXLOG - 1 ; lg >= 0 ; --lg)
    {
        int next = parent[u][lg];
        if(next != 0 && !isAnc(next, v))
        {
            u = next;
        }
    }

    return parent[u][0];
}

void initDFS(int node, int par, int dp)
{
    depth[node] = dp;
    in[node] = ++timer;

    for(int lg = 0 ; lg < MAXLOG ; ++lg)
    {
        if(lg == 0)
        {
            parent[node][lg] = par;
            continue;
        }

        int prev = parent[node][lg - 1];
        parent[node][lg] = parent[prev][lg - 1];
    }
    
    for(const int &to : g[node])
    {
        if(to == par)
        {
            continue;
        }

        initDFS(to, node, dp + 1);
    }

    out[node] = timer;
}

std::vector<int> getPath(int start, int end)
{
    int lca = getLCA(start, end);

    std::vector<int> prefix;
    while(start != lca)
    {
        prefix.push_back(start);
        start = parent[start][0];
    }

    std::vector<int> suffix;
    while(end != lca)
    {
        suffix.push_back(end);
        end = parent[end][0];
    }

    std::reverse(suffix.begin(), suffix.end());

    std::vector<int> path;
    for(int x : prefix) path.push_back(x);
    path.push_back(lca);
    for(int x : suffix) path.push_back(x);
    return path;
}

void markDown(int row, int col)
{
    for(int r = row ; r < MAXN ; ++r)
    {
        if(position[r][col] != 0)
        {
            isBad[r][col] = 1;
        }
        else
        {
            break;
        }
    }
}

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;
    k = D.size();

    for(auto &[u, v] : R)
    {
        g[u + 1].push_back(v + 1);
        g[v + 1].push_back(u + 1);
    }

    initDFS(1, 0, 1);

    for(int i = 1 ; i <= k ; ++i)
    {
        std::vector<int> p = getPath(S[i - 1] + 1, T[i - 1] + 1);
        pathLen[i] = p.size();
        for(int j = D[i - 1], ptr = 0; j < D[i - 1] + p.size() ; ++j, ++ptr)
        {
            position[j][i] = p[ptr];
            isBad[j][i] = I[i - 1];
        }
    }

    for(int j = 0; j < MAXN; ++j)
    {
        for(int first = 1 ; first <= k ; ++first)
        {
            for(int second = 1 ; second <= k ; ++second)
            {
                if(isBad[j][first] && position[j][first] == position[j][second])
                {
                    markDown(j, second);
                }
            }
        }
    }

    ans.resize(k, 0);
    for(int i = 1 ; i <= k ; ++i)
    {
        int first = -1;
        for(int j = 0; j < MAXN; ++j)
        {
            if(first == -1 && isBad[j][i])
            {
                first = j;
                break;
            }
        }

        ans[i - 1] = (D[i - 1] + pathLen[i] - 1) - first + 1;
    }

    return ans;
}
