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


const int MAX_N = 5000;

vector <int> susNa[MAX_N];

int tim = 1;
int checked[MAX_N];
int from[MAX_N];
void dfs (int start, int endd, vector<int>& path) {
    queue<int> toCheck;
    toCheck.push(start);
    while(!toCheck.empty()) {
        int curr = toCheck.front();
        toCheck.pop();
        if (curr == endd) break; 
        if (checked[curr] == tim) continue;
        checked[curr] = tim;

        for (int i = 0; i < susNa[curr].size(); i ++) {
            int sus = susNa[curr][i];
            if (checked[sus] == tim) continue;
            from[sus] = curr;
            toCheck.push(sus);
        }
    }
    int curr = endd;
    while (curr != start) {
        path.push_back(curr);
        curr = from[curr];
    }
    path.push_back(start);
    reverse(path.begin(), path.end());
    tim ++;
}

vector<int> solve(vector<pair<int, int>> R, vector<long long> D, vector<int> S, vector<int> T, vector<bool> I) {
    for (int i = 0; i < MAX_N; i ++) susNa[i].clear();
    for (int i = 0; i < R.size(); i ++) {
        auto curr = R[i];
        susNa[curr.first].push_back(curr.second);
        susNa[curr.second].push_back(curr.first);
    }

    int M = S.size();
    // cout << M << endl;
    
    vector <int> paths[MAX_N];
    
    unsigned long long maxx = 0;

    for (int i = 0; i < M; i ++) {
        dfs(S[i], T[i], paths[i]);
        maxx = max(maxx, paths[i].size() + D[i] + 1);
        
    }

    // for (int i = 0; i < M; i ++) {
    //     for (int x = 0; x < D[i]; x ++) cout << "  ";
    //     for (int j = 0; j < paths[i].size(); j ++) {
    //         cout << paths[i][j] << " ";
    //     }
    //     cout << endl;
    // }

    
    
    bool isInfected[MAX_N];
    vector <int> ans;
    for (int i = 0; i < M; i ++) {
        ans.push_back(0);
        isInfected[i] = false;
    }

    vector <int> infected;
    for (int i = 0; i < M; i ++) {
        if (I[i]) infected.push_back(i);
        isInfected[i] = I[i];
    }

    for (int j = 0; j <= maxx; j ++) {
        // cout << "At day " << j << endl;
        for (int i = 0; i < infected.size(); i ++) {
            int curr = infected[i]; 
            int dayForCurrent = j - D[curr];
            if (dayForCurrent < 0 || dayForCurrent >= paths[curr].size()) continue;
            ans[curr] ++;
            // cout << "Adding 1 to " << curr << endl;
            for (int x = 0; x < M; x++) {
                int dayForX = j - D[x];
                if (isInfected[x]) continue;
                if (dayForX >= paths[x].size() || dayForX < 0) continue;
                if (paths[x][dayForX] == paths[curr][dayForCurrent]) {
                    // cout << x << " gets infected by " << curr << endl;
                    isInfected[x] = true;
                    infected.push_back(x);
                }            
            }
        }
    }
    return ans;
}
