[이코테] 그래프 알고리즘 - 여행 계획

해당 포스팅의 문제의 출처는 나동빈님의 이것이 취업을 위한 코딩 테스트 교재를 공부하면서 풀이할 때 본인의 사고 과정 흐름과 문제 풀이를 기록하기 위함 입니다. 

 

---

문제설명

한울이가 사는 나라에는 N개의 여행지가 있으며, 각 여행지는 1 ~ N번까지의 번호로 구분된다. 또한 임의의 두 여행지 사이에는 두 여행지를 연결하는 도로가 존재할 수 있다. 이 때, 여행지가 도로로 연결되어 있다면 양방향으로 이동이 가능하다는 의미이다. 한울이는 하나의 여행 계획을 세운 뒤에 이 여행 계획이 가능한지의 여부를 판단하고자 한다. 예를 들어 N = 5 이고, 다음과 같이 도로의 정보가 주어졌다고 가정하자.

 

• 1번 여행지 - 2번 여행지
• 1번 여행지 - 4번 여행지
• 1번 여행지 - 5번 여행지
• 2번 여행지 - 3번 여행지
• 2번 여행지 - 4번 여행지

만약 한울이의 여행 계획이 2번 -> 3번 -> 4번 -> 3번 이라고 하자. 이 경우, 2번 -> 3번 -> 2번 -> 4번 -> 2번 -> 3번의 순서로 여행지를 방문하면 여행 계획을 따를 수 있다. 여행지의 개수와 여행지 간의 연결 정보가 주어졌을 때, 한울이의 여행 계획이 가능한지의 여부를 판별하는 프로그램을 작성해라.

입력조건

• 첫째 줄에 여행자의 수 N과 여행 계획에 속한 도시의 수 M이 주어진다.(1 <= N, M <= 500)
• 다음 N개의 줄에 걸쳐 N X N 행렬을 통해 임의의 두 여행지가 서로 연결되어 있는지의 여부가 주어진다. 그 값이 1이라면 서로 연결되었다는 의미이며, 0이라면 서로 연결되어 있지 않다는 의미이다.
• 마지막 줄에 한울이의 여행 계획에 포함된 여행지의 번호들이 주어지며 공백으로 구분된다.

출력조건

• 첫째 줄에 한울이의 여행 계획이 가능하다면 YES를, 불가능하다면 NO를 출력해라.

사고과정

• 문제에서 "N개의 여행지, 도로, 양방향" 키워드를 뽑아내어 그래프 알고리즘을 판단했다. 그래프 알고리즘 중 최소 신장 트리를 찾는 크루스칼인지, 위상 정렬인지 문제와 대응시켜보다가 문제에서 "여행 계획 여부 판단"이라는 키워드를 통해 주어진 노드들 간에 서로소 집합 여부를 판단하는 문제라는 것을 힘겹게 캐치했다. 아이디어를 얻고 난 후 구현에 들어가서 시간 내에 풀이를 완료했다.
• 힘들었던 또 다른 점은 간선 연결 정보를 2차원 행렬로 받아서 이를 (a, b) 형태로 어떻게 바꾸어 놓을까에 대해서도 시간을 어느정도 썼다.
• 다만 책 풀이와 약간 다른 점은 필자 코드보다 리팩토링을 매우 잘해놓았다. 필자는 간선 정보 입력을 다 받은 후에 union 연산을 처리했지만 책 속에서는 N X N 행렬의 1행씩 받을 때마다 간선 정보로 바꾸어서 union 연산을 수행한 점이 인상적이었다. 또 필자는 union 연산을 모두 수행하고 부모테이블에서 입력받은 한울이의 여행 계획 원소마다 find 연산을 수행한 후 해당 부모 테이블 값이 모두 같으면(set으로 감싸주고 길이가 1이면) YES를 출력하도록 처리했으나, 책 속에서는 이것도 바로 한울이의 여행 계획 원소를 m-1번 만큼 반복해서 앞, 뒤로 find 연산을 취해 부모노드가 같은지 처리하는 식으로 깔끔하게 해놓았다. 끝이 없는 리팩토링..

풀이

1. 책 속 풀이를 권장하므로 책 풀이를 먼저 제시한다.

import sys

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


n, m = map(int, input().split())
# 부모 테이블 초기화
parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

for i in range(n):
    data = list(map(int, input().split()))
    for j in range(n):
        # 연결된 간선이 있는 노드 2개에 대해 union 연산 수행
        if data[j] == 1:
            union_parent(parent, i+1, j+1)

plans = list(map(int, input().split()))
# 여행 계획에 있는 도시들의 루트노드가 같은지 확인
result = True
for i in range(m-1):
    if find_parent(parent, plans[i]) != find_parent(parent, plans[i+1]):
        result = False

if result:
    print('YES')
else:
    print('NO')

2. 내 풀이

n, m = map(int, input().split())
# 주어지는 N by N 행렬을 받아서 (a, b) 형태로 바꾸기
array = []
for _ in range(n):
    array.append(list(map(int, input().split())))
edges = []
for i in range(n):
    for j in range(n):
        if array[i][j] == 1:
            edges.append((i+1, j+1))
# 한울이의 여행 계획 입력받기
plans = list(map(int, input().split()))
# 여행지의 부모 테이블 초기화
parent = [0] * (n+1)
for i in range(1, n+1):
    parent[i] = i

# find, Union 연산
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

# 주어진 연결 여부 정보 하나씩 union 연산
for i in range(len(edges)):
    a, b = edges[i]
    union_parent(parent, a, b)
# 각 원소의 루트노드 탐색
for i in range(1, n+1):
    find_parent(parent, i)

# 여행 계획에 포함된 여행지들의 루트노드가 모두 같다면 여행 가능
root_node = set([parent[p] for p in plans])
if len(root_node) == 1:
    print('YES')
else:
    print('NO')