#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;
    int cnt = 10;
    while (curr != start && cnt) {
        path.push_back(curr);
        curr = from[curr];
        cnt --;
    }
    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];
    
    int maxx = -1, maxx2 = -1;

    for (int i = 0; i < M; i ++) {
        dfs(S[i], T[i], paths[i]);
        int sz = paths[i].size();
        if (sz > maxx) {
            maxx2 = maxx;
            maxx = sz;
        } else if (sz > maxx2) {
            maxx2 = sz;
        }
    }

    // for (int i = 0; i < M; i ++) {
    //     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 ++) {
        for (int i = 0; i < infected.size(); i ++) {
            int curr = infected[i]; 
            ans[curr] ++;

            for (int x = 0; x < M; x++) {
                if (isInfected[x]) continue;
                if (j >= paths[x].size()) continue;
                if (paths[x][j] == paths[curr][j]) {
                    isInfected[x] = true;
                    infected.push_back(x);
                }            
            }
        }
    }
    return ans;
}
