이것이 코딩테스트다 41 여행 계획 (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이면 연결되어 있지 않다는 의미
마지막줄에는 여행 계획이 포함된 여행지의 번호들이 주어진다.

---

Test Case

입력
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

---

문제 풀이

1. 서로소의 판별(union-find)로 쉽게 풀 수 있는 문제인데, 이 문제 처럼 연결정보를 그래프로 주어질 수도 있고, (노드, 노드) 이런식으로 줄 수도 있다. 주어진 정보를 내가 아는 코드로 표현할 줄 알아야 한다. 

2. 서로소 판별할 때 부모테이블이 꼭 필요하다는 점을 기억해야 한다.

3. 부모 노드가 같은지 판별하고 부모가 같으면 계속 이어나가고 다르면 바로 끝내면 되는 부분에서 flag 값을 줬다.

---

Source Code

import sys
input = sys.stdin.readline

# 노드수 n, 여행계획 수 m
n,m = map(int, input().split())

parent = [0] *(n+1)
for i in range(1,n+1):
    parent[i] = i

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

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

plan = list(map(int, input().split()))

result = False

for i in range(m-1):
    if find_parent(parent, plan[i]) != find_parent(parent,plan[i-1]):
        result = True
        break

if result == True:
    print("NO")
else :
    print("YES")