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

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;
    for (long long i = 0; i < M; i++) {
        lastDay = std::max(lastDay, D[i] + abs(T[i] - S[i]) + 1);
    }

    std::vector<long long> currentCity(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];
                continue;
            }

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

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

            if (T[i] > currentCity[i]) {
                currentCity[i]++;
            } else {
                currentCity[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;
}
