백준 1976번: 여행 가자 [python]

난이도 : 골드 4
유형 : 유니온 파인드
https://www.acmicpc.net/problem/1976

문제 접근

• 여러 도시가 있고, 어떤 도시는 서로 도로로 연결되어 있다.
• 여행 계획이 주어졌을 때, 모든 도시를 순서대로 방문할 수 있는가?를 판단하는 문제.
• 핵심은 각 도시가 같은 연결된 집합(=같은 그룹, 같은 부모) 에 속해 있는지를 판별하는 것

핵심 아이디어

• 도시 간의 직접 연결 정보가 인접행렬로 주어진다.
• 연결 여부를 바탕으로 유니온-파인드(Disjoint Set Union, DSU)를 활용해 각 도시의 대표 부모를 병합한다.
• 여행 계획에 포함된 모든 도시가 같은 부모(루트) 를 가지고 있으면 YES, 아니면 NO.

해답 및 풀이

import sys
input = sys.stdin.readline

# 도시 수, 여행 계획 도시 수
n = int(input())
m = int(input())

# 유니온 파인드를 위한 부모 배열
parent = [i for i in range(n + 1)]

# find 연산 - 경로 압축 포함
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

# union 연산
def union(a, b):
    root_a = find(a)
    root_b = find(b)
    if root_a != root_b:
        parent[root_b] = root_a

# 인접행렬 입력 받아 union 처리
for i in range(1, n + 1):
    row = list(map(int, input().split()))
    for j in range(1, n + 1):
        if row[j - 1] == 1:
            union(i, j)

# 여행 계획
plan = list(map(int, input().split()))
root = find(plan[0])

# 모든 도시가 같은 집합(루트)이면 YES
for city in plan[1:]:
    if find(city) != root:
        print("NO")
        break
else:
    print("YES")

유니온 파인드는 연결 요소를 빠르게 관리할 수 있는 좋은 방법이다.
여행 경로에 있는 도시들이 전부 같은 대표 도시를 갖고 있는지만 판단하면 된다