[이코테] 여행 계획

문제

 한울이가 사는 나라에는 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

---

풀이

도로는 양방향으로 연결되어 있기 때문에 무향그래프이며, 여행 계획이 가능한지 여부는 계획한 여행지가 모두 같은 집합에 속하는 지 여부로 판별할 수 있습니다. 따라서 서로소 집합 알고리즘을 활용하면 쉽게 풀 수 있는 문제입니다.

 

링크: https://jcy1996.tistory.com/34?category=960192 

서로소 집합 알고리즘

 

전체 코드는 다음과 같습니다.

n, m = list(map(int, input().split()))
graph = []
for _ in range(n):
    graph.append(list(map(int, input().split())))
target = list(map(int, input().split()))

parent = [0] * n
for i in range(n):
    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):
    for j in range(i + 1, n):
        if graph[i][j] == 1:
            union_parent(parent, i, j)
        
base = parent[target[0]]
for t in list(set(target)):
    t -= 1
    if parent[t] != base:
        print('NO')
        break
else:
    print('YES')

 같은 집합에 속하는지 여부를 판별하는 것이 관건이기 때문에 경로 압축을 사용합니다. 미리 주어진 인접행렬 정보를 통해 만약 두 노드가 연결되어 있다면 Union연산을 행합니다. 무향 그래프의 인접행렬은 대칭적이기 때문에, 첫번째 노드가 i이라면 두 번째 노드는 i 이후의 노드만 고려하게 합니다.

for i in range(n):
    for j in range(i + 1, n):
        if graph[i][j] == 1:
            union_parent(parent, i, j)

 

 결과를 출력하는 부분은 간단합니다. 기준점을 가장 첫 여행지의 루트노드로 잡고, 이후 모든 여행지에 대해서 첫 여행지의 루트노드와 같은지, 즉 같은 집합에 속하는지 판별합니다. 하나라도 예외가 있다면 'NO'를 출력하고 순회를 멈춥니다. 만약 브레이크 체커에 걸리지 않는다면 'YES'를 반환합니다.

 단 중복되는 노드에 대해서는 다시 고려할 필요가 없기 때문에 집합으로 한번 형변환을 하여 중복값을 삭제한 뒤 다시 리스트로 형변환해주었습니다.

base = parent[target[0]]
for t in list(set(target)):
    t -= 1
    if parent[t] != base:
        print('NO')
        break
else:
    print('YES')