[#알고리즘/Union-Find/[백준]1976번: 여행 가자(python)]

Question

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

Solution

앞서 풀었던 union-find 알고리즘 문제와 동작이 유사하다. 다만 유의해야할 점은 예제 입력만을 보고 도시의 수와 여행 계획에 속한 도시의 수가 3개로 정해진 것이 아니므로 리스트 등의 자료구조로 입력을 받아와야 한다.

Code

import sys

input = sys.stdin.readline

# find
def find_parent(parent, x) :
    if parent[x] != x :
        parent[x] = find_parent(parent, parent[x])
    return parent[x]     

# union
def union_parent(parent, a, b) :
    a = find_parent(parent, a)
    b = find_parent(parent, b)
    if a < b :
        parent[b] = a
    else :
        parent[a] = b

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

# 부모 테이블 생성 및 초기화
parent = [0] * (n + 1)

for i in range(1, n + 1) :
    parent[i] = i
    
for i in range(1, n + 1) :
    city = list(map(int, input().split()))
    for j in range(n) :
        if city[j] == 1 :
            union_parent(parent, i, j + 1)

trip_city = map(int, input().split())
result = set(find_parent(parent, i) for i in trip_city)

if len(result) != 1 :
    print('NO')
else :
    print('YES')