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

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
) {
    int N = R.size() + 1;
    std::vector <std::vector <int>> bfs;

    for (int i = 0; i < N; i++) {
        std::vector <int> t;
        for (int j = 0; j < N; j++) {
            t.push_back(-1);
        }
        bfs.push_back(t);
    }

    std::vector <std::vector <int>> edges(N);
    for (auto i: R) {
        edges[i.first].push_back(i.second);
        edges[i.second].push_back(i.first);
    }

    for (int s = 0; s < N; s++) {
        bfs[s][s] = 0;
        std::queue <int> q;
        q.push(s);

        while (!q.empty()) {
            int t = q.front(); q.pop();
            for (int i: edges[t]) {
                if (bfs[s][i] == -1) {
                    bfs[s][i] = bfs[s][t]+1;
                    q.push(i);
                }
            }
        }
    }

    std::vector <int> p(S.size());
    std::vector <bool> inf(S.size());
    std::vector <int> res(S.size());
    for (int i = 0; i < S.size(); i++) {
        p[i] = S[i];
        inf[i] = I[i];
    }

    std::vector <bool> terminated(S.size());
    int t_sum = 0;

    for (int t = 0; t < 1000; t++) {
        if (t_sum == S.size()) {
            break;
        }

        std::vector <std::vector <int>> l(N);
        for (int i = 0; i < p.size(); i++) {
            l[p[i]].push_back(i);
        }

        for (int i = 0; i < p.size(); i++) {
            for (int j: l[i]) {
                for (int k: l[i]) {
                    if (j == k) continue;
                    if (inf[j]) inf[k] = true;
                    if (inf[k]) inf[j] = true;
                }
            }
        }

        for (int i = 0; i < S.size(); i++) {
            res[i] += inf[i];
        }

        //updating positions

        for (int ind = 0; ind < p.size(); ind++) {
            int i = p[ind];
            if (terminated[i]) {
                continue;
            }

            int best_dist = 1000;
            int best = -1;
            for (int j: edges[i]) {
                if (bfs[j][T[i]] < best_dist) {
                    best_dist = bfs[j][T[i]];
                    best = j;
                }
            }
            if (best_dist == 0) {
                terminated.push_back(i);
                t_sum++;
            } else {
                p[ind] = best;
            }
        }
    }

    return res;
}
