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

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

 

# 백준 1976 여행 가자
# Disjoint Set (Union-Find) 알고리즘 활용

import sys
sys.setrecursionlimit(10**7)

N = int(input())  # 도시 개수
M = int(input())  # 여행에 포함된 도시 개수

# 부모 배열 초기화 (처음엔 자기 자신을 부모로 설정)
parent = [i for i in range(N + 1)]


# find 연산 (경로 압축)
def find(x):
    if x != parent[x]:
        parent[x] = find(parent[x])  # 루트까지 올라가면서 부모 갱신
    return parent[x]


# union 연산 (두 집합 합치기)
def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[a] = b  # 한쪽 루트를 다른 쪽 밑에 붙임


# 두 도시가 같은 집합(연결된 네트워크)에 속하는지 확인
def is_con(a, b):
    return find(a) == find(b)


# 인접 행렬 입력 처리 (도시 간 연결 여부)
for A in range(1, N + 1):
    can_go = list(map(int, input().split()))  # A번째 도시에서 다른 도시들로 가는 경로
    for B in range(1, N + 1):
        if can_go[B - 1] == 1:   # A와 B가 연결되어 있다면
            if A > B:            # (중복 union 방지: 대각선 위쪽만 확인)
                union(A, B)


# 여행 경로 입력
routes = list(map(int, input().split()))

# 여행 가능 여부 검사
succ_flag = True
for i in range(M - 1):
    # 연속된 도시들이 같은 네트워크(집합)에 속해 있어야 여행 가능
    if not is_con(routes[i], routes[i + 1]):
        succ_flag = False

# 결과 출력
if succ_flag:
    print('YES')
else:
    print('NO')

UNION-FIND 정복해보자..