[백준] #1976. 여행 가자 (python)

문제

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

1976번: 여행 가자

동혁이는 친구들과 함께 여행을 가려고 한다. 한국에는 도시가 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라는 여행경로를 통해 목적을 달성할 수 있다.

 

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

 

예시 입력

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

 

예시 출력

YES

 

풀이

 

이 문제는 유니온 파인드(Union-find) 알고리즘을 사용하면 된다.

 

import sys
from collections import deque

n=int(sys.stdin.readline())
m=int(sys.stdin.readline())
city=list(range(201))
possible=True

def find(node):
    if city[node]==node:
        return node
    else:
        city[node]=find(city[node])
        return city[node]

def union(a, b):
    a=find(a)
    b=find(b)
    city[b]=city[a]

#(1) - 도시의 연결 정보를 기반으로 도시 연결
for i in range(1, n+1):
    nums=list(map(int, sys.stdin.readline().split()))
    for idx, num in enumerate(nums):
        if num==1:
            union(i, idx+1)

plans=deque(list(map(int, sys.stdin.readline().split())))
while len(plans)>1:
    first=plans.popleft()
    second=plans.popleft()
    a=find(first)
    b=find(second)
    if a!=b:
        possible=False
        break
    plans.appendleft(second)

if possible:
    print('YES')
else:
    print('NO')

문제에서 도시의 연결 정보를 제공하기 때문에 (1)에서 union 연산을 사용해 도시끼리 연결을 진행한다.