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

void bfs(const std::vector<std::vector<long long>> &graph, std::vector<bool> &visited, std::vector<long long> &distance, std::vector<long long> &parent, long long node) {
    std::queue<long long> q;
    q.push(node);
    distance[node] = 0;

    while (!q.empty()) {
        node = q.front();
        q.pop();

        if (visited[node]) {
            continue;
        }
        visited[node] = true;

        for (const long long &next : graph[node]) {
            if (distance[next] > distance[node] + 1) {
                distance[next] = distance[node] + 1;
                parent[next] = node;
                q.push(next);
            }
        }
    }
}

std::vector<int> solve(
    std::vector<std::pair<int, int>> R, //edges
    std::vector<long long> D, //start day
    std::vector<int> S, //start city
    std::vector<int> T, //final city
    std::vector<bool> I //infected
) 
{
    const long long M = D.size();
    const long long N = R.size() + 1;

    long long lastDay = -1;
    std::vector<std::vector<long long>> graph(N + 1, std::vector<long long>());
    std::vector<std::vector<long long>> shortestPaths(M, std::vector<long long>());

    for (const auto &p : R) {
        graph[p.first].push_back(p.second);
        graph[p.second].push_back(p.first);
    }

    for (long long i = 0; i < M; i++) {
        std::vector<bool> visited(N + 1, false);
        std::vector<long long> distance(N + 1, LONG_LONG_MAX);
        std::vector<long long> parent(N + 1, -1);

        bfs(graph, visited, distance, parent, S[i]);
    
        long long current = T[i];
        while (current != -1) {
            shortestPaths[i].push_back(current);
            current = parent[current];
        }
        std::reverse(shortestPaths[i].begin(), shortestPaths[i].end());

        lastDay = std::max(lastDay, D[i] + (long long)shortestPaths[i].size() + 1);   
    }

    std::vector<long long> currentCity(M, -1);
    std::vector<long long> currentStep(M, -1);
    std::vector<int> answer(M, 0);

    for (long long currentDay = 0; currentDay <= lastDay; currentDay++) {
        for (long long i = 0; i < M; i++) {
            if (currentDay == D[i]) {
                currentCity[i] = S[i];
                currentStep[i] = 0;
                continue;
            }

            if (currentCity[i] == -1) {
                continue;
            }

            if (currentCity[i] == T[i]) {
                currentCity[i] = -1;
                continue;
            }

            currentStep[i] = currentStep[i] + 1;
            currentCity[i] = shortestPaths[i][currentStep[i]];
        }

        std::vector<bool> cityInfected(N + 1, false);

        for (long long i = 0; i < M; i++) {
            if (!I[i] or currentCity[i] == -1) {
                continue;
            }
            cityInfected[currentCity[i]] = true;
        }

        for (long long i = 0; i < M; i++) {
            if (currentCity[i] != -1 and cityInfected[currentCity[i]]) {
                I[i] = true;
                answer[i]++;
            }
        }
    }

    return answer;
}
