1081D - Maximum Distance - CodeForces Solution

#include <bits/extc++.h>
using namespace std; 

#define all(x) begin(x), end(x)
#define rep(i, a, b) for(int i = a; i < b; i ++) 
typedef vector<int> vi;
typedef pair<int, int> pii;
#define fir first
#define sec second
#define pb push_back

struct E {
    int w, u, v;
};

vi fa;
int find(int p) {
    if(fa[p] == p) return p;
    return fa[p] = find(fa[p]);
}

vi dep, sz;

void dfs(int p, int q, vector<bool>& v, vector<vector<pii>>& g) {
    if(q != -1) dep[p] = dep[q] + 1;
    sz[p] = 0;
    if(v[p]) sz[p] = 1;
    for(auto [nxt, w] : g[p]) {
        if(nxt == q) continue ;
        dfs(nxt, p, v, g);
        sz[p] += sz[nxt];
    }
}

int main() {
    int n, m, k; cin >> n >> m >> k;
    vector<bool> v(n, 0);
    rep(i, 0, k) {
        int ai; cin >> ai; --ai;
        v[ai] = 1;
    }
    vector<vector<pii>> e(n), g(n);
    fa = vi(n, 0);
    iota(all(fa), 0);
    vector<E> es, e2;
    rep(i, 0, m) {
        int u, v, w; cin >> u >> v >> w; --u, --v;
        e[u].pb({v, w}), e[v].pb({u, w});
        es.pb({w, u, v});
    }
    sort(all(es), [&](E x, E y) {return x.w < y.w;});
    for(auto [w, u, v] : es) {
        int gu = find(u), gv = find(v);
        if(gu == gv) continue ; 
        fa[gu] = gv; 
        g[u].pb({v, w});
        g[v].pb({u, w});
        e2.pb({w, u, v});
    }
    reverse(all(e2)); 

    dep = sz = vi(n, 0);
    dfs(0, -1, v, g);
    int ans = 0;
    for(auto [w, u, v] : e2) {
        int x = u; if(dep[v] > dep[x]) x = v;
        if(sz[x] == 0 || sz[x] == k) continue ;
        else {ans = w; break ;}
    }
    rep(i, 0, k) cout << ans << " ";
    cout << endl;
}