골드4

문제

해당 문제 링크

정답

## 여행 가자

import sys
sys.setrecursionlimit(10**6)

input=sys.stdin.readline

def find(x):
    if parent[x]!=x:
        parent[x]=find(parent[x])
    return parent[x]

def union(a, b):
    a=find(a)
    b=find(b)

    if a<b:
        parent[b]=a
    else:
        parent[a]=b

#유니온 파인드
n=int(input())
m=int(input())

parent=[i for i in range(n+1)]

for i in range(1,n+1):
    edge = list(map(int, input().split()))
    for j in range(0, n):
        if edge[j]==1:
            union(i, j+1)

# parent = [-1] + parent
# print(parent)

path=list(map(int, input().split()))

# start=parent[path[0]]
for i in range(1, m):
     # 맞지 않다면 같은 집합에 속해져 있지 않기 때문에 NO라는 문구가 뜸.
    if parent[path[i]] != parent[path[0]]:
        print("NO")
        break
else:
    print("YES")

풀이

참고한 다른 사람 코드에서 내가 수정해서 작성해보았다.

참고한 코드는 유니온 파인드의 값을 넣을 때 1부터가 아닌 0부터 넣고 경로 체크를 할때는 -1를 넣어서 parents 배열 자체가 -1, 0, 0, 0으로 만들어지게 한다.
즉 예시에 맞게 보자면, union(0, 1), union(1, 0), union(1, 2), union(2, 1)으로 들어가게 된다는 것이다. 1, 2, 3 경로를 지나가는 것은 리스트의 인덱스로 확인하기 때문에 문제가 없다.
코드가 틀리지는 않았지만 난 좀 더 수정하였다. 코드를 작성할 때 헷갈리지 않는다면 참고한 코드처럼 작성하는게 더 간단하다!

내가 수정한 코드는 원래 유니온 파인드에서 하는 것과 같이 union(1, 2), union(2, 1), union(2, 3), union(3, 2)으로 들어가게 된다. 그리하여 parent배열은 0, 1, 1, 1이 된다. 이때 따로 -1을 안붙여도 되는 이유는 union 연산을 할 때 0을 건들지 않았기 때문이다.