여행 계획 코드 및 해설 (파이썬)

출처: 이것이 취업을 위한 코딩 테스트다 with 파이썬, Q 41, p 393

 

서로소 집합을 이용해 여행 계획 안에 있는 여행지들이 모두 같은 집합에 속하면 예행 계획이 가능한 것으로, 그렇지 않으면 불가능한 것으로 풀이했습니다.

연결된 도시들의 정보가 리스트가 아닌 행렬로 주어지므로 행렬의 값 중 0이 아니라 1인 부분의 좌표를 통해 연결된 도시들의 정보를 담은 리스트 links를 생성했습니다. 이를 이용해 서로소 집합 알고리즘을 수행하고, 계획에 있는 도시들이 속한 집합을 확인해 이들이 모두 같으면 YES, 아니면 NO를 출력했습니다.

 

import sys

# 입력 받아오기 
inputs = list(map(int, sys.stdin.read().split()))
# 여행지의 수
N = inputs[0]
# 여행 계획 속 여행지의 개수 
M = inputs[1]
# 연결 정보
graph = inputs[2:(N*N)+2]
links = []
for i, g in enumerate(graph):
    if g == 1:
        # 상삼각행렬 내 좌표로 변환 
        link = sorted((i//5+1, i%5+1))
        if link not in links:
            links.append(link)

# 여행 계획
# 여행 계획 내 여행지가 같은 집합에 있으면 가능 
plans = inputs[(N*N)+2:]

# x가 속한 집합 찾기 
def find_parent(parent, x):
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])
    return parent[x]

# 집합 합치기 
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

# 부모 테이블을 자기 자신으로 초기화 
parent = [0] * (N+1)
for i in range(1, N+1):
    parent[i] = i

# 연결되어 있으면 같은 집합으로 합치기 
for a, b in links:
    union_parent(parent, a, b)

# 계획에 있는 도시들이 속한 집합들 
plan_parents = []
for plan in plans:
    plan_parents.append(find_parent(parent, plan))

# 계획에 있는 모든 도시들이 같은 집합에 속해있으면 성공 
if len(set(plan_parents)) == 1:
    print('YES')
else:
    print('NO')