#include "infection.h"
#include <bits/stdc++.h>

using integer = int;
using namespace std;

#define int long long

struct om
{
    int D, S, T;
    bool infectat;
};

int n, m;
vector<int> g[150005], f[150005];
om v[150005];
int h[150005], sz[150005], t[150005];
int bl[20][150005];

void dfs(int nod)
{
    sz[nod] = 1;
    for (auto vecin : g[nod])
    {
        if (vecin != t[nod])
        {
            t[vecin] = nod;
            h[vecin] = 1 + h[nod];
            f[nod].push_back(vecin);
            dfs(vecin);
            sz[nod] += sz[vecin];
        }
    }
}

int lca(int x, int y)
{
    if (h[x] < h[y])
        swap(x, y);
    for (int pas = 18; pas >= 0; pas--)
        if (h[x] - (1 << pas) >= h[y])
            x = bl[pas][x];
    if (x == y)
        return x;
    for (int pas = 18; pas >= 0; pas--)
        if (bl[pas][x] != bl[pas][y])
            x = bl[pas][x], y = bl[pas][y];
    return t[x];
}

int dist(int x, int y)
{
    return h[x] + h[y] - 2 * h[lca(x, y)];
}

vector<int> dr[5005];
vector<integer> ans;

vector<int> drum(int x, int y)
{
    int lc = lca(x, y);
    vector<int> pth;
    pth.push_back(x);
    while (x != lc)
    {
        x = t[x];
        pth.push_back(x);
    }
    vector<int> p2;
    while (y != lc)
    {
        p2.push_back(y);
        y = t[y];
    }
    reverse(p2.begin(), p2.end());
    for (auto it : p2)
        pth.push_back(it);
    return pth;
}

bool blk[5005][5005];

vector<integer> solve(vector<pair<integer, integer>> R, vector<int> D, vector<integer> S, vector<integer> T, vector<bool> I)
{
    n = 1 + R.size();
    m = D.size();
    for (auto it : R)
    {
        g[it.first + 1].push_back(it.second + 1);
        g[it.second + 1].push_back(it.first + 1);
    }
    for (int i = 0; i < m; i++)
    {
        v[i].D = D[i];
        v[i].S = S[i] + 1;
        v[i].T = T[i] + 1;
        v[i].infectat = I[i];
    }
    dfs(1);
    for (int i = 1; i <= n; i++)
        bl[0][i] = t[i];
    for (int j = 1; j <= 18; j++)
        for (int i = 1; i <= n; i++)
            bl[j][i] = bl[j - 1][bl[j - 1][i]];
    for (int i = 0; i < m; i++)
        dr[i] = drum(v[i].S, v[i].T);
    ans.resize(m);
    for (int i = 0; i < m; i++)
        if (v[i].infectat)
            ans[i] = dr[i].size();
    for (int x = 0; x <= n; x++)
    {
        for (int i = 0; i < m; i++)
        {
            if (dr[i].size() > x)
            {
                if (v[i].infectat == true)
                {
                    blk[x][dr[i][x]] = true;
                }
            }
        }
        for (int i = 0; i < m; i++)
        {
            if (dr[i].size() > x)
            {
                if (!v[i].infectat)
                {
                    if (blk[x][dr[i][x]])
                    {
                        ans[i] = dr[i].size() - x;
                        v[i].infectat = true;
                    }
                }
            }
        }
    }
    return ans;
}
