[BOJ/Gold 2] 백준 9344 도로(C++)

문제 링크

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

9344번: 도로

 

문제

어떤 나라에는 1부터 N까지 이름 붙여진 N개의 도시가 있다. 한 엔지니어는 모든 도시를 연결하는 도로를 건설하고자 한다. 즉, 모든 도시에 대해 항상 다른 어떤 도시로든 이동할 수 있어야 한다. 이때 여러 개의 도시를 통과할 수도 있다. 그의 팀은 몇 개의 길(도로 후보)을 조사했다. 각각의 길은 두 도시를 양방향으로 잇는다. 길 위에 도로를 지을 때는 특정 비용이 든다. (길이 짧을수록 비용도 싸다.)

이 엔지니어는 교통 시스템을 미리 계획하지 않았다. 그는 그저 선호에 따라 한 개의 길을 선택하고, 도로를 건설하는 일을 모든 도시가 연결될 때까지 반복한다.

지금 엔지니어는 도시 p와 도시 q를 잇는 도로를 건설하고자 한다. 비용을 감축하라는 정부의 압력에 의해, 그는 당신에게 그가 해당 도로를 지어야 하는지 여부를 판단하는 프로그램을 작성할 것을 요구했다. 당신의 프로그램은 그 도로를 지으면서 모든 도시를 연결하는 가장 짧은 도로망을 만들 수 있으면 YES라고 대답해야 한다. 그렇지 않다면, NO를 출력해야 한다.

입력

첫 줄에 테스트 케이스의 개수 T가 주어진다. (T ≤ 10) 

각 테스트 케이스의 첫 줄에는 4개의 정수 N, M, p, q가 주어진다. N(2 ≤ N ≤ 10,000)은 도로망 위에 존재하는 도시의 수이다. M(1 ≤ M ≤ 20,000)은 길의 수이다. p와 q(1 ≤ p,q ≤ N)는 그 사이에 도로를 지어도 되는지 판단해야 하는 두 도시이다.

이어지는 M개의 줄 각각에는 u, v, w가 주어진다.(1 ≤ u ≤ N, 1 ≤ v ≤ N, 1 ≤ w ≤ 400,000) 도시 u와 v를 잇는 양방향 길의 비용이 w라는 것을 의미한다. 도로를 짓는 데 드는 비용은 모두 다르며, 두 도시를 잇는 길은 오직 하나이다. 모든 도시를 잇는 도로망이 최소 한 개 이상 존재한다는 것이 보장된다. 모든 입력은 정수이다.

출력

각 테스트 케이스에 대해, p-q를 지으면서 가장 짧은 도로망을 만들 수 있으면 YES를 출력한다. 아니면 NO를 출력한다.

예제 입력 1

3
2 1 1 2
1 2 4
3 3 2 3
1 2 1
1 3 2
2 3 3
4 5 3 4
1 2 1
1 3 3
3 4 2
1 4 4
4 2 5

예제 출력 1

YES
NO
YES

알고리즘 분류

• 크루스칼 알고리즘

풀이

각 테스트 케이스마다 최소 스패닝 트리를 만들고, P와 Q를 잇는 간선이 포함되는 지 확인하고 포함된다면 YES, 포함되지 않는다면 NO를 출력한다.

코드

#include <iostream>
#include <string>
#include <cmath>
#include <set>
#include <map>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#include <climits>
#define FASTIO cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define MAX 10001
#define LL long long
#define INF 1e18

using namespace std;
struct INFO {
	int Start, End;
	LL Cost;
};

int T, N, M, P, Q;
vector<INFO> Edge;
int Parent[MAX];
bool Flag;
LL answer;

void Init() {
	for (int i = 1; i <= N; i++) {
		Parent[i] = i;
	}
	Edge.clear();
	Flag = false;
	answer = 0;
}

int Find(int X) {
	if (Parent[X] == X) {
		return X;
	}
	return Parent[X] = Find(Parent[X]);
}

void Union(int X, int Y) {
	int PX = Find(X);
	int PY = Find(Y);
	if (PX < PY) {
		Parent[PY] = PX;
	}
	else {
		Parent[PX] = PY;
	}
}

bool Comp(INFO A, INFO B) {
	return (A.Cost < B.Cost);
}

void Settings() {
	sort(Edge.begin(), Edge.end(), Comp);
	for (int i = 0; i < Edge.size(); i++) {
		int U = Edge[i].Start;
		int V = Edge[i].End;
		LL C = Edge[i].Cost;
		if (Find(U) != Find(V)) {
			Union(U, V);
			answer += C;
			if (((P == U) && (Q == V)) || ((P == V) && (Q == U))) {
				Flag = true;
			}
		}
	}
}

void Find_Answer() {
	if (Flag) {
		cout << "YES\n";
	}
	else {
		cout << "NO\n";
	}
}

void Input() {
	cin >> T;
	while (T--) {
		cin >> N >> M >> P >> Q;
		Init();
		for (int i = 0; i < M; i++) {
			int U, V;
			LL W;
			cin >> U >> V >> W;
			Edge.push_back({ U,V,W });
		}
		Settings();
		Find_Answer();
	};
}

int main() {
	FASTIO
	Input();

	return 0;
}