[백준] 1976 - 여행 가자 (파이썬)

n = int(input())
m = int(input())
array = [-1 for _ in range(n+1)]
def find(x):
    if array[x] < 0:
        return x
    array[x] = find(array[x])
    return array[x]

def union(x,y):
    x = find(x)
    y = find(y)
    if x==y:
        return
    array[x] = y

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

answer = list(map(int,input().split()))
if m == 0 or m==1:
    print('YES')
else:
    t = 0
    for i in range(len(answer)-1):
        if find(answer[i]) != find(answer[i+1]):
            t = 1
    if t==0:
        print('YES')
    else:
        print("NO")

유니온 파인드 알고리즘을 사용하여 나타내었다.

t = 0 과 t = 1을 나눠 연결되었는지 확인하였다.