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

using namespace std;

int dist[5005][5005];
vector<int>adj[5005];

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;
    int m = I.size();
    vector<int>ans(m,0);

    for(int i=0; i<n-1; i++) {
        adj[R[i].first].push_back(R[i].second);
        adj[R[i].second].push_back(R[i].first);
    }

    for(int i=0; i<n; i++) {
        queue<pair<int,int>>q;
        vector<int>visited(n+1,0);
        q.push({i,0});
        while(!q.empty()) {
            pair<int,int> v = q.front(); q.pop();
            if(visited[v.first]) continue;
            visited[v.first] = 1;
            dist[i][v.first] = v.second;
            for(auto u : adj[v.first]) {
                if(!visited[u]) q.push({u, v.second+1});
            }
        }
    }
    
    for(int i=0; i<m; i++) {
        if(I[i]) ans[i] = dist[S[i]][T[i]]+1;
    }

    // simulacija
    int grad[m]; // trenutni grad partygoera i
    vector<long long>inf[n];
    for(int i=0; i<m; i++) grad[i] = S[i];
    priority_queue<pair<int,int>>pq;
    for(int i=0; i<m; i++) {
        pq.push({I[i], i});
    }
    while(!pq.empty()) { 
        int tren = pq.top().second; // partygoer kojeg pomjeramo
        //cout << "tren-" << tren << " tmp-" << -tmp << endl;
        int s = grad[tren];
        //cout << s << endl;
        for(int x : inf[s]) {
            if(x == D[tren]+dist[S[tren]][s]) {
                if(!I[tren]) {
                    I[tren] = 1;
                    pq.push({I[tren], tren});
                    //cout << "tren-" << tren << " T[tren]-" << T[tren] << " dist-" << dist[tren][T[tren]] << endl;
                    ans[tren] = dist[s][T[tren]]+1;
                }
            }
        }
        if(I[tren]) {
            //cout << "tren-" << tren << " s-" << s << endl;
            inf[s].push_back(D[tren]+dist[S[tren]][s]);
        }
        if(grad[tren] == T[tren]) pq.pop();
        else {
            for(auto u : adj[s]) {
                if(dist[u][T[tren]] < dist[s][T[tren]]) {
                    grad[tren] = u;
                    break;
                }
            }
        }
    }
    return ans;
}