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

const int MAX = 4 * 5002;
int niv[MAX + 2], turIn[MAX + 2], par[MAX + 2];
int rmq[20][MAX + 2], lg[MAX + 2];
deque<int> tras[MAX + 2];
vector<int> gr[MAX + 2];

struct Moment {
    long long om, timp;
};
unordered_map<int, vector<Moment>> travers;

static inline void DFS(int nod, int tata = 0) {
    niv[nod] = 1 + niv[tata];
    rmq[0][++turIn[0]] = nod;
    turIn[nod] = turIn[0];
    par[nod] = tata;
    for(int fiu : gr[nod]) {
        if(fiu != tata) {
            DFS(fiu, nod);
            rmq[0][++turIn[0]] = nod;
        }
    }
}

static inline int MinNiv(int a, int b) {
    return (niv[a] < niv[b] ? a : b);
}

static inline void RMQ() {
    for(int i = 2; i <= turIn[0]; i++) lg[i] = 1 + lg[i >> 1];
    for(int i = 1; i <= lg[turIn[0]]; i++) {
        for(int j = 1; j + (1 << i) - 1 <= turIn[0]; j++) {
            rmq[i][j] = MinNiv(
                rmq[i - 1][j],
                rmq[i - 1][j + (1 << (i - 1))]
            );
        }
    }
}

static inline int LCA(int a, int b) {
    a = turIn[a];
    b = turIn[b];
    if(a > b) swap(a, b);
    int k = lg[b - a];
    return MinNiv(rmq[k][a], rmq[k][b - (1 << k) + 1]);
}

vector<int> solve(
    vector<pair<int, int>> r,
    vector<long long> d,
    vector<int> s,
    vector<int> t,
    vector<bool> infectat
) {
    for(int i = 0; i < r.size(); i++) {
        gr[r[i].first].push_back(r[i].second);
        gr[r[i].second].push_back(r[i].first);
    }
    //cout << '\n';

    DFS(0);
    RMQ();

    vector<int> nopti(d.size(), 0);
    for(int i = 0; i < d.size(); i++) {
        vector<int> drum;

        int st = s[i];
        int dr = t[i];
        int lca = LCA(st, dr);

        while(lca != st) {
            tras[i].push_back(st);
            st = par[st];
        }
        tras[i].push_back(lca);
        while(lca != dr) {
            drum.push_back(dr);
            dr = par[dr];
        }
        reverse(drum.begin(), drum.end());
        for(int cur : drum) tras[i].push_back(cur);

        if(infectat[i]) nopti[i] = tras[i].size();

        long long tp = d[i];
        for(int cur : tras[i]) {
            travers[cur].push_back({i, tp});
            tp++;
        }
    }

    map<long long, map<int, vector<int>>> tpev;
    for(int i = 0; i < r.size() + 1; i++) {
        for(Moment mom : travers[i]) {
            tpev[mom.timp][i].push_back(mom.om);
        }
    }

    for(auto& timp : tpev) {
        for(auto& nod : timp.second) {
            bool inf = false;

            for(int om : nod.second) inf |= infectat[om];
            for(int om : nod.second) {
                if(inf && !infectat[om]) {
                    nopti[om] = tras[om].size() - (timp.first - d[om]);
                    infectat[om] = true;
                }
            }
        }
    }

    return nopti;
}


/*
9 4
0 1
1 2
2 3
3 4
4 5
3 6
6 7
6 8
0 0 4 0
0 5 8 1
0 3 5 0
3 7 8 0
*/

