[백준 1976 파이썬] 여행 가자 ( 골드 4, 유니온 파인드)

알고리즘 유형 : 유니온 파인드
풀이 참고 없이 스스로 풀었나요? : O

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

소스 코드(파이썬)

import sys
input = sys.stdin.readline

# 유니온 파인드

def find(x):
    if graph[x] < 0:
        return x
        
    graph[x] = find(graph[x])
    return graph[x]
    
def union(x, y):
    x = find(x)
    y = find(y)
    
    if x == y:
        return
    
    if graph[x] < graph[y]:
        graph[y] = x
    elif graph[y] < graph[x]:
        graph[x] = y
    else:
        graph[x] = y
        graph[y] -= 1
        
    return

N = int(input())
M = int(input())
graph = [-1]*(N+1)

# 연결 관계가 1인 두 노드의 루트 노드를 합쳐준다.(그래프 합치기)
for start in range(1, N+1):
    for end, isLinked in zip(range(1, N+1), map(int, input().split())):
        if isLinked:
            union(start, end)
   
path = list(set(map(int, input().split())))
result = "YES"
# 방문 경로 내 도시 중, 한 쌍이라도 루트 노드가 서로 다르다면
# 즉, 서로가 한 그래프 내에 존재하지 않는다면 방문 경로는
# 절대로 다 돌 수 없으므로 result = "NO"
# 한 편, len(path)가 1인 경우는 그냥 그 도시만 방문하면되니까
# YES 출력
for i in range(1, len(path)):
    if find(path[0]) != find(path[i]):
        result = "NO"
        break
print(result)

풀이 요약

1. 이 문제는 유니온 파인드를 적용하여 풀면 된다.

   이 풀이는 weighted union find를 활용한 풀이이다.

2. 그래프 리스트를 -1로 초기화하고, 연결 관계 입력값을 순회하면서, 값이 1일 때의 두 노드를 union해준다.

3. 이렇게 완성된 그래프에 대해, 입력받은 방문 경로를 순회하면서, 이 경로 중 한 쌍이라도 서로 루트 노드가 다른 도시가 있다면, 서로 방문 불가능한 도시가 발생한 것이므로 result에 NO를 저장하고 break 후 result를 출력해준다.

   그러한 케이스가 하나도 발생하지 않은 경우(=모든 도시가 한 트리 내에 존재함), 그리고 방문 경로가 도시 하나뿐인 경우에는 YES가 출력된다.

배운 점, 어려웠던 점

• 유니온 파인드 활용력을 기를 수 있어 좋았다.