백준 1976번 여행 가자 / C++ / 유니온 파인드(Union-Find)

알고리즘 개념: 유니온-파인드

유니온 파인드에 대해 이전에 정리해 둔 개념 정리를 한번 읽어보고 오는 것을 추천합니다.

[알고리즘] 유니온 파인드 (Union-Find) 개념과 최적화 기법

[알고리즘] 유니온 파인드 (Union-Find) 개념과 최적화 기법

 

백준 1976번 여행 가자

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

 

예제 보려면 더보기 클릭

예제 입력 1

3
3
0 1 0
1 0 1
0 1 0
1 2 3

 

예제 출력 1

YES

---

문제풀이

 i번째 줄의 j번째 수는 i번 도시와 j번 도시의 연결 정보를 의미한다.

처음엔 이게 무슨 말인가 싶었는데 예제와 함께 보니 이해가 됐다.

예제에서 첫번째 0 1 0은 1번 도시가 2와 연결되어 있다는 것이고

두번째 1 0 1은 2번 도시가 1번과 3번에 연결되어 있다는 것이다. 

 

그리고 우리가 주목해야 할점은 하나 더 있다. 

1이면 연결된 것이고 0이면 연결이 되지 않은 것이다.

즉 0일땐 굳이 숫자를 볼 필요가 없다는 것이다. 

나는 이중 for문을 통해 하나의 변수로 입력을 받고 1일때 연결을 해주도록 코드로 구현했다. 

for (int i = 1; i <= n; i++) {
	for (int j = 1; j <= n; j++) {
		cin >> a;
		if (a == 1)
			merge(i, j);
	}
}

 

 

Union과 Find를 구현한 함수는 Union-Find 개념에 설명이 올라와 있다. 

Find 연산 구현

int find(int n) {
	if (p[n] == n) return n;
	else return p[n] = find(p[n]);
}

Union 연산 구현

void merge(int a, int b){
    a = find(a);
    b = find(b);
    if (a==b) return;
    else p[b] = a;
 }

 

그리고 이제 마지막으로 주목해야 할 점은

마지막 줄에는 여행 계획이 주어진다. 

이걸 어떻게 구현해야 하나 고민이 많았는데 먼저 처음으로 주는 값을 start 즉 여행의 시작지로 받고

그 후로 받는 값들을 직접 find(n)함수로 여행을 떠나보면서 비교해보는 것으로 코드를 구현했다. 

int start;
int cnt = 0;
cin >> start;
for (int i = 1; i < m; i++) {
	cin >> a;
	if (find(start) == find(a)) {
		cnt++;
	}
	else break;
	start = a;
}

그렇다면 cnt의 개수가 m-1과 같다면! 모든 도시를 여행한 것이 된다.

(하나라도 여행을 못떠나면 break문이 실행되므로 cnt값이 늘지 않는다)

if (cnt == m - 1) cout << "YES";
else cout << "NO";

이것을 모두 정리하여 하나로 합치면 아래 코드가 된다. 

---

제출한 코드(C++)

#include <iostream>
#include <string>
#include <algorithm>
#include <vector>
#define INF 1000005
using namespace std;
using ll = long long int;
int p[INF] = { 0 };
int n, m;
void fillp(int n) {
	for (int i = 0; i <= n; i++) 
		p[i] = i;
}
int find(int n) {
	if (p[n] == n) return n;
	else return p[n] = find(p[n]);
}
void merge(int a, int b) {
	a = find(a);
	b = find(b);
	if (a == b) return;
	p[b] = a;
}
int main(void) {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	cin >> n >> m;
	fillp(n);
	int a;
	for (int i = 1; i <= n; i++) {
		for (int j = 1; j <= n; j++) {
			cin >> a;
			if (a == 1)
				merge(i, j);
		}
	}
	int start;
	int cnt = 0;
	cin >> start;
	for (int i = 1; i < m; i++) {
		cin >> a;
		if (find(start) == find(a)) {
			cnt++;
		}
		else break;
		start = a;
	}
	if (cnt == m - 1) cout << "YES";
	else cout << "NO";

	return 0;
}