#include "infection.h"
// #include <vector>
// #include <utility>

#include <bits/stdc++.h>

using namespace std;

int n, m;
vector<int> g[5005];
int p[5005][5005];

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
) {
    // fac un bfs din fiecare om ca sa imi fac un arbore rootat in el
    // fiecare om o sa faca maxim 5000 de pasi deci pot sa simulez efectiv?
    // dau skip cat timp nu am pe nimeni
    n = R.size() + 1;
    m = D.size();

    vector<int> ans(m, 0);

    for(auto e : R) {
        g[e.first].push_back(e.second);
        g[e.second].push_back(e.first);
    }
    for(int i = 0; i < n; i++) {
        for(int j = 0; j < n; j++) p[i][j] = -2;
        p[i][i] = -1;
        queue<int> q;
        q.push(i);
        while(!q.empty()) {
            auto node = q.front();
            q.pop();
            for(auto nxt : g[node]) {
                if(p[i][nxt] == -2) {
                    p[i][nxt] = node;
                    q.push(nxt);
                }
            }
        }
    }

    // for(int i = 0; i < n; i++) {
    //     cout << i << '\n';
    //     for(int j = 0; j < n; j++) {
    //         cout << p[i][j] << ' ';
    //     }
    //     cout << '\n';
    // }

    // merg prin fiecare om
    // ii stabilesc pathul 
    // marchez ca in ziua x omul se afla acolo
    // fac un map pt asta?
    unordered_map<long long, unordered_map<int, vector<int>>> mp;
    for(int i = 0; i < m; i++) {
        // plec din S[i] si trebuie sa ajung in T[i]
        // la ziua D[i] + timer voi fi pe un nod. pot sa il marchez

        // cout << i << ' ' << S[i] << ' ' << T[i] << '\n';

        vector<int> path;
        int node = T[i];
        while(node != -1) {
            path.push_back(node);
            node = p[S[i]][node];
        }
        reverse(path.begin(), path.end());

        // cout << "Pathul de la " << S[i] << " la " << T[i] << " pentru " << i << "este:\n";
        // for(int timer = 0; timer < path.size(); timer++) {
        //     cout << path[timer] << ' ';
        // }
        // cout << '\n';

        for(int timer = 0; timer < path.size(); timer++) {
            mp[D[i] + timer][path[timer]].push_back(i);
        }

    }

    for(auto elem : mp) {
        // elem.fs e timpul curent
        // elem.sd are un map in care pt nodul i stiu cine e in el

        // cout << "Ziua: " << elem.first << '\n';

        for(auto node : elem.second) {
            // node.fs e nodul curent
            // verific daca am vreun infected

            // cout << "Nodul: " << node.first << '\n';

            bool infected = false;
            for(auto who : node.second) {

                // cout << who << ' ';

                if(I[who] == 1) {
                    infected = true;
                    break;
                }
            }

            // cout << '\n';

            for(auto who : node.second) {
                if(infected) I[who] = 1;
                if(I[who] == 1) {
                    ans[who]++;
                }
            }
        }

        // cout << '\n';

    }

    return ans;
}
