[백준1976번] 여행 가자 / Python3

문제

동혁이는 친구들과 함께 여행을 가려고 한다. 한국에는 도시가 N개 있고 임의의 두 도시 사이에 길이 있을 수도, 없을 수도 있다. 동혁이의 여행 일정이 주어졌을 때, 이 여행 경로가 가능한 것인지 알아보자. 물론 중간에 다른 도시를 경유해서 여행을 할 수도 있다. 예를 들어 도시가 5개 있고, A-B, B-C, A-D, B-D, E-A의 길이 있고, 동혁이의 여행 계획이 E C B C D 라면 E-A-B-C-B-C-B-D라는 여행경로를 통해 목적을 달성할 수 있다.

도시들의 개수와 도시들 간의 연결 여부가 주어져 있고, 동혁이의 여행 계획에 속한 도시들이 순서대로 주어졌을 때(중복 가능) 가능한지 여부를 판별하는 프로그램을 작성하시오.

입력

첫 줄에 도시의 수 N이 주어진다. N은 200이하이다. 둘째 줄에 여행 계획에 속한 도시들의 수 M이 주어진다. M은 1000이하이다. 다음 N * N 행렬을 통해 임의의 두 도시가 연결되었는지에 관한 정보가 주어진다. 1이면 연결된 것이고 0이면 연결이 되지 않은 것이다. A와 B가 연결되었으면 B와 A도 연결되어 있다. 마지막 줄에는 여행 계획이 주어진다.

출력

첫 줄에 가능하면 YES 불가능하면 NO를 출력한다.

예제 입력 1

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

예제 출력 1

YES

나의 풀이

얼핏 이 문제를 접하면 BFS나 DFS로 풀면 될 것 같다. 그러나 중복이 가능하기 때문에 유니온파인드를 사용하면 쉽게 풀 수 있다. 아마도 BFS를 사용한다면 중복을 가능하게 해야하므로 메모리초과가 발생하지 않을까 싶다. 무튼 유니온 파인드에 대한 설명은 백준1717번 해당 문제에서 이미 설명을 해놓았다.

무튼 이 문제에 대한 풀이를 해보면 우선 입력받은 인접행렬을 가지고 union으로 서로소 집합을 만들어주어야 한다. 그래프를 입력받고 1 (즉, 연결이 되어있는 지점) 에 대하여 각각의 행과 열을 union으로 집합을 합쳐주면 된다. 물론 여기서 그래프의 모든 원소를 확인할 필요는 없다. 왜냐, 이 문제에서 A와 B가 연결되어 있다면 B와 A도 연결이 되어있다고 알려주었으니 입력되는 인접행렬은 대칭행렬이라는 의미이고 이는 곧 대각선을 기준으로 아래 삼각형부분만 확인해도 된다는 말이기때문이다. 따라서 아래의 코드에서 n을 i로만 바꾸어주면 된다.

for i in range(n):
    for j in range(n):        # n을 i로 바꾸면 대각선 아래 삼각형만 확인 가능
        if graph[i][j] == 1:
            union(i+1,j+1)

인접행렬을 가지고 서로소 집합들을 나누었다면 이제 여행경로라고 입력을 받은 곳들이 모두 한 집합에 속하는지를 확인해주면 문제를 해결할 수 있게된다. 필자는 여행경로를 리스트로 입력받아 find()함수를 사용해 모든 여행령로의 root가 같은지를 확인해서 결과값을 출력해주었다.

코드

# 1976번 여행 가자
import sys

# find
def find(x):
    global parents

    while True:
        if parents[x] == x:
            return x
        x = parents[x]

# union
def union(a,b):
    global parents

    a = find(a)
    b = find(b)

    if not a == b:
        parents[b] = a

# main
n = int(input())
m = int(input())

# 인접행렬 입력
graph = []
for _ in range(n):
    graph.append([int(x) for x in sys.stdin.readline().split()])

# 서로소 집합 
parents = [int(x) for x in range(n+1)]
for i in range(n):
    for j in range(n):
        if graph[i][j] == 1:
            union(i+1,j+1)

# 여행 경로 입력
travel = [int(x) for x in sys.stdin.readline().split()]

# 여행 경로에 대한 모든 root 확인
root = find(travel[0])
isDis = True
for t in travel:
    if not root == find(t):
        isDis = False
        break

print("YES" if isDis else "NO")