[백준] 1976-여행가자 (유니온 파인드)

김영민·2024년 9월 8일
0

코딩테스트

목록 보기
28/32



코드

    
N = int(input())
M = int(input())

connect = [list(map(int,input().split(" "))) for _ in range(N)]

graph = [[] for _ in range(N+1)]

def find_root(x):
    if nodes[x] != x:
        nodes[x] = find_root(nodes[x])
    return nodes[x]

def union_root(a, b):
    a = find_root(a)
    b = find_root(b)
    if a < b:
        nodes[b] = a
    else:
        nodes[a] = b

for a in range(N):
    for b in range(N):
        if connect[a][b] == 1:
            graph[a+1].append(b+1)


plan = list(map(int,input().split(" ")))

nodes = [i for i in range(N+1)]

for i in range(len(graph)):
    for num in graph[i]:
        union_root(i,num)


start = nodes[plan[0]]
flag = 0
for i in range(1,M):
    if nodes[plan[i]] != start:
        print("NO")
        flag = 1
        break

if flag == 0:
    print("YES")

풀이

  • 갔던 도시를 다시 가도 된다는 점에서 visited list가 필요없다고 생각.
  • 조상이 같으면 어떻게든 갈 수 있으므로, union-find로 풀이.
  • 조상 찾기 진행 (union-find)
  • 첫번째 도시를 조상으로 설정하고, 각 도시들의 조상 도시가 첫번째 도시가 아니면 NO를 출력

0개의 댓글