[BOJ] 9344번 : 도로 - C++

$N$개의 도시를 가장 적은 비용으로 이어야 합니다. 즉, MST를 구하는 문제입니다. 그런데 특정 도시 $p, q$를 연결하는 도로를 사용하면서 MST를 만들 수 있는지를 판단하는 문제입니다. 단순하게 접근하면 일반적인 MST를 수행하면서 p, q를 잇는 도로가 사용되는지를 판단하면 되는데, 만약 모든 간선의 비용이 같다고 생각해 봅시다. p, q 도로를 사용하지 않고도 같은 비용으로 MST를 만들 수 있습니다. 따라서 pq 도로와 길이가 같은 도로들 중에서 pq 도로가 우선적으로 사용되도록 해야 합니다. priority queue의 정렬 함수를 적절하게 바꿔 정렬하면 구현할 수 있습니다.

 

조심해야 할 부분은 priority queue의 정렬 함수는 sort에 사용하는 정렬 함수랑 반대라는 것입니다.

 

#include "bits/stdc++.h"

using namespace std;

int parent[10001];

int find(int x) {
    if (parent[x] == x)
        return x;
    else
        return parent[x] = find(parent[x]);
}

void unite(int a, int b) {
    int pa = find(a);
    int pb = find(b);

    if (pa != pb) {
        parent[pa] = pb;
    }
}

int p, q;

struct Cmp {
    bool operator()(const pair<int, pair<int, int>> &a, const pair<int, pair<int, int>> &b) {
        if (a.first == b.first) {
            if (
                a.second.first == p && a.second.second == q ||
                a.second.first == q && a.second.second == p
            )
                return false;
            else if (
                    b.second.first == p && b.second.second == q ||
                    b.second.first == q && b.second.second == p
            )
                return true;
        }

        return a > b;
    }
};

int main() {
    cin.tie(nullptr);
    cout.tie(nullptr);
    ios_base::sync_with_stdio(false);

    int t;
    cin >> t;
    while (t--) {
        int n, m;
        cin >> n >> m >> p >> q;
        for (int i = 0; i <= n; i++)
            parent[i] = i;

        priority_queue<
                pair<int, pair<int, int>>,
                vector<pair<int, pair<int, int>>>,
                Cmp
        > pq;

        while (m--) {
            int a, b, c;
            cin >> a >> b >> c;
            pq.push({ c, { a, b }});
        }

        bool found = false;
        while (!pq.empty()) {
            auto [cost, pp] = pq.top();
            auto [a, b] = pp;
            pq.pop();

            if (find(a) != find(b)) {
                unite(a, b);

                if (a == p && b == q || a == q && b == p) {
                    found = true;
                    break;
                }
            }
        }

        if (found)
            cout << "YES\n";
        else
            cout << "NO\n";
    }
}