[백준] 1976번 여행 가자 파이썬 코드

코드

# 최단거리 X, 연결되어있는지 확인
# 집합에 속해 있는지 확인하는 과정
# 유니온-파인드 이용

import sys
sys.setrecursionlimit(10**5)
input = sys.stdin.readline

n = int(input())
m = int(input())
node = [i for i in range(n+1)] # 부모 노드

def find(a, parent):
  if a != parent[a]: # 루트 노드가 아니면
    parent[a] = find(parent[a], parent) # 루트 노드 찾기
  return parent[a]

def union(a,b,parent):
  a = find(a, parent)
  b = find(b, parent)
  if a < b: # 루트 노드가 작은 값에 합친다
    parent[b] = a
  else:
    parent[a] = b

for i in range(1,n+1):
  lst = [*map(int, input().split())]
  for j in range(len(lst)):
    if lst[j] == 1:
      union(i,j+1,node) # 그래프 만들어주기
      
travel = [*map(int, input().split())] # 여행지

result = set() # 중복을 없애기 위해 집합 이용

for i in travel:
  result.add(find(i, node)) # find 함수를 이용해 루트 노드 찾아주기
print("YES") if len(result) == 1 else print("NO") # 1개가 아니면 여행 불가

루트 노드는 parent 배열에서 보장되지 않을 수 있으므로, 재귀를 통해 루트 노드를 찾아주어야 한다.