[Graph] 이코테 “여행 계획” Python 풀이

문제

• 한울이가 사는 나라에는 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를 출력합니다.

입력 예시

5 4
0 1 0 1 1
1 0 1 1 0
0 1 0 0 0
1 1 0 0 0
1 0 0 0 0
2 3 4 3

출력 예시

YES

풀이

• 이 문제는 서로소 집합 알고리즘을 이용하여, 그래프에서 노드 간의 연결성을 파악해 해결할 수 있다.
• 우리가 여기에서 알 수 있는 점은, ‘여행 계획’에 해당하는 모든 노드가 같은 집합에 속하기만 하면, 가능한 여행 경로라는 것이다.
• 여행 경로를 출력하는게 아닌, 여행 경로가 가능한지만 판단하면 되기 때문이다.
• 사이클은 판단할 필요가 없다.

예제 코드

import sys

# 여행지가 속한 집합 찾기
def find(parent, x):
    # 해당 여행지가 루트 집합이 아니라면, 루트 집합 찾기
    if (parent[x] != x):
        parent[x] = find(parent, parent[x])
    return parent[x]

# 두 여행지가 속한 집합을 합치기
def union(parent, x, y):
    a = find(parent, x)
    b = find(parent, y)
    if (a < b):
        parent[b] = a
    else:
        parent[a] = b

# 여행지의 수 n과 여행 계획에 속한 도시의 수 m
n, m = map(int, sys.stdin.readline().split())
# 각 여행지가 속한 집합을 저장하기 위한 1차원 배열
parent = [i for i in range(n+1)]

# 여행지 연결 정보
for i in range(n):
    row = list(map(int, sys.stdin.readline().split()))
    for j in range(n):
        # 해당 여행지가 연결되어 있다면
        if (row[j] == 1):
            # 간선 정보에 따라 합집합 진행
            union(parent, i, j)

# 여행지 계획
plan = list(map(int, sys.stdin.readline().split()))

# 여행 계획에 속하는 모든 노드의 루트가 동일한지 확인
result = True
for i in range(m - 1):
    if find(parent, plan[i]) != find(parent, plan[i + 1]):
        result = False
        break

# 여행 경로의 노드가 같은 모두 집합인지 확인
if (result == True):
    print('YES')
else:
    print('NO')

참고

• 이것이 취업을 위한 코딩 테스트다. with python