[백준 1976번] 여행 가자 (C++, Python)

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

1976번: 여행 가자

이번 포스트에서는 백준 알고리즘 문제인 여행 가자를 포스트하겠습니다.

 

해당 문제는 solved.ac 기준 골드 4에 위치한 문제입니다.

 

먼저 문제입니다.

입력으로 도시의 수 N, 여행 계획에 속하는 도시의 수 M, 그리고, 도시의 연결 정보를 입력받고, 여행 계획을 입력받습니다.

 

이때, 입력받은 도시의 연결 정보를 바탕으로 여행 계획이 가능한지 여부를 출력하는 문제입니다.

 

해당 문제는 그래프 문제로, dfs, bfs, 플로이드 워셜 등으로 해결이 가능합니다.

 

필자는 해당 문제를 유니온 파인드로 해결한 케이스를 설명하겠습니다.

 

유니온 파인드 알고리즘은 여러 개의 노드를 공통된 루트 노드로 묶고 다른 노드의 루트 노드가 같은지 비교하는 알고리즘입니다. (쉽게 같은 집합에 속해있는지 검사)

 

해당 문제에서는 연결 정보를 받고, 이어져 있는 도시들은 같은 집합으로 묶어 주어서, 입력받은 여행 루트를 검사할 때, 집합이 같은지 다른지를 검사하여 여행 루트가 가능한지 확인할 수 있습니다.

 

먼저 C++ 코드입니다.

 

#include <iostream>
 
using namespace std;
 
int city[201];
int res[1001] = {0,};
 
int Find(int key) {
    if(key == city[key])    return key;
    else return city[key] = find(city[key]);
}
 
void Union(int a, int b) {
    int ta = Find(a);
    int tb = Find(b);
    if(ta != tb)    city[ta] = tb;
}
 
int main() {
    int n, m, tmp;
    cin >> n >> m;
    
    for(int i = 1; i <= 200; i++)
        city[i] = i;
 
    for(int i = 1; i <= n; i++) {
        for(int k = i; k <= n; k++) {
            cin >> tmp;
            if(tmp)
                Union(i,k);
        }
    }
 
    int prev = 0;
    for(int i = 1; i <= m; i++) {
        cin >> tmp;
        res[i] = Find(tmp);
        if(i == 1)    prev = Find(tmp);
        else {
            if(res[i] != prev) {
                cout << "NO";
                return 0;
            }
        }
    }
    
    cout << "YES";
 
    return 0;
}

다음은 Python 코드입니다.

 

city = [0] * 201
res = [0] * 1001
 
def find(key) :
    if key == city[key] :
        return key
    else :
        city[key] = find(city[key])
        return city[key]
 
def union(a, b) :
    ta = find(a)
    tb = find(b)
 
    if ta != tb :
        city[ta] = tb
 
n = int(input())
m = int(input())
 
for i in range(1, 201) :
    city[i] = i
 
for i in range(1, n + 1) :
    arr = list(map(int, input().split()))
    for k in range(1, n + 1) :
        if arr[k - 1] == 1 :
            union(i, k)
 
prev = 0
tmp = list(map(int, input().split()))
for i in range(1, m + 1) :
    res[i] = find(tmp[i - 1])
    if i == 1 :
        prev = find(tmp[i - 1])
    else :
        if res[i] != prev :
            print("NO")
            exit()
print("YES")

유니온 파인드 알고리즘에서 먼저 Find 함수는 해당 노드가 속해있는 집합의 루트 노드를 찾는 함수입니다.

(key와 city[key]가 다를경우, city[key]로 재탐색하여 city[key]와 key가 같아질때 까지 탐색 -> city[key]와 key가 같으면 루트노드)

EX) 1 -> 2, 2-> 3, 3->4, 4->5 는 5라는 루트 노드 아래 4가 있고, 그 아래 1,2,3 이 속해있다고 정의할 수 있음.

 

위 처럼 루트노드로 묶지 않으면, 1->5 로 가는 길이 가능한지 확인하기 위해 많은 노드를 방문해야하므로, 하나의 루트노드로 묶어 줍니다.

 

이 처럼 루트노드로 묶는 것은 Union 함수에서 처리합니다. a에서 b로 가는 길이 있을 때, a의 루트 노드아래 b의 루트 노드를 두는 작업을 처리합니다.

 

이때, Union 함수에서 루트노드를 찾기 위해 Find함수를 호출하게 되는데, 해당 값이 루트가 아닐 경우, 재귀적으로 탐색하여 루트 노드를 찾고 그 값으로 이전에 호출한 Find에서 가리키는 값이 가리키는 노드를 탐색한 루트 노드로 변경하여 하나의 집합으로 묶어줍니다. ( 위 코드의 Find 함수에서 else : return city[key] = find(city[key]) 의 부분입니다. )

 

이렇게 연결정보를 가지고 Union 함수를 통해 하나의 집합으로 묶어줬으므로, 가능한 경로라면 그 경로는 루트 노드가 같게 나옵니다. 따라서, 문제에서 요구한 여행 경로가 가능한지 판단 가능합니다.

해당 문제는 유니온 파인드 알고리즘을 알고 있다면, 문제에서 유니온 파인드가 필요한 문제임을 빠르게 찾아낼 수 있는 문제였습니다. 또한 N이 200까지 주어지고 시간제한이 2초이므로, O(N^3)인 플로이드 워셜 알고리즘이나 BFS, DFS로도 해결할 수 있는 문제이니 다른 알고리즘으로도 해결해 보는것도 재밌을거같은 문제입니다.

 

다음 포스트에서도 다른 백준 알고리즘으로 포스트 해보겠습니다. 감사합니다.