5542번: JOI 국가의 행사 - 재검토 필요

https://www.acmicpc.net/problem/5542

#include<stdio.h>
#include<algorithm>
#include<vector>
#include<queue>
using namespace std;
const int MAX_N = 1e6, MAX_Q = 1e6;
int n, m, k, q, flag[MAX_N + 1], res[MAX_Q];
bool ck[MAX_N + 1];
vector<pair<int, int> > adj[MAX_N + 1], v, query[MAX_N + 1];
priority_queue<pair<int, int> > pq;
vector<int> st[MAX_N + 1];
int f(int h) {
    if (h == flag[h]) return h;
    return flag[h] = f(flag[h]);
}
int main() {
    scanf("%d %d %d %d", &n, &m, &k, &q);
    for (int i = 0; i < m; i++) {
        int a, b, c;
        scanf("%d %d %d", &a, &b, &c);
        adj[a].push_back({ b, c });
        adj[b].push_back({ a, c });
    }
    for (int i = 0; i < k; i++) {
        int a;
        scanf("%d", &a);
        pq.push({ 0, a });
    }
    while (!pq.empty()) {
        int h = pq.top().second, dis = -pq.top().first;
        pq.pop();
        if (ck[h]) continue;
        ck[h] = true;
        v.push_back({ h, dis });
        for (auto t : adj[h]) pq.push({ -dis - t.second, t.first });
    }
    for (int i = 0; i < q; i++) {
        int a, b;
        scanf("%d %d", &a, &b);
        query[a].push_back({ b, i });
        query[b].push_back({ a, i });
    }
    for (int i = 1; i <= n; i++) flag[i] = i, st[i].push_back(i);
    fill(res, res + q, -1);
    for (int i = v.size() - 1; i >= 0; i--) {
        pair<int, int> h = v[i];
        ck[h.first] = false;
        for (auto t : adj[h.first]) {
            if (!ck[t.first]) {
                int a = f(h.first), b = f(t.first);
                if (a == b) continue;
                if (st[a].size() < st[b].size()) swap(a, b);
                flag[b] = a;
                for (auto sit : st[b]) {
                    st[a].push_back(sit);
                    for (auto it : query[sit])
                        if (res[it.second] == -1 && f(sit) == f(it.first)) res[it.second] = h.second;
                }
            }
        }
    }
    for (int i = 0; i < q; i++) printf("%d\n", res[i]);
    return 0;
}