[백준 1976] 여행가자

문제

N개의 도시가 있다. 특정 두 도시사이에는 길이 있을 수도, 없을 수도 있다. 이때 여행하려고하는 도시의 경로가 주어졌을 때 해당 경로가 가능한지 알아보자.

같은 도시를 두 번 이상 방문하는 것도 가능하다.

 

입력

첫째 줄에 도시의 수 N을 입력하고, 이어서 여행 계획에 속한 도시의 수 M을 입력한다.

그 다음 줄부터 N개의 줄에 도시의 연결정보를 2차원 리스트로 표시하고 연결이 안된 것은 0, 된 것은 1로 표시한다.

가장 마지막 줄에는 여행 경로를 나타내는 도시 M개를 입력한다.

 

출력

입력된 여행 경로가 가능하면 YES, 불가능하면 NO를 출력한다.

---

접근

만약 A->B->C의 여행경로가 있다고 했을 때 B와 C가 A에 속해져있으면 이 여행경로로 가능하고, 다른 그룹에 속해있다면 불가능한 것이다. 따라서 도시들의 연결정보를 입력받았을 때 각각의 도시에 대한 대표노드를 구하고 여행경로 순서대로 비교하여 대표노드가 모두 같다면 YES, 다르다면 NO를 출력한다.

 

구현

- 입력받아야 될 것을 모두 입력받은 뒤에

- 도시들이 속해있는 대표노드를 찾는다. city_list의 초기에는 자기자신을 가리키고 있다.

- 이어서 2차원 리스트를 탐색하여 j도시를 i도시에 합친다.

city_list = [i for i in range(N)]
for i in range(N):
    for j in range(N):
        if graph[i][j]:
            union(i, j)

 

- 이제 여행 경로에있는 도시대로 여행할 수 있는지 판단을 해야한다.

- 이를 위해 이전도시와 현재도시를 계속해서 비교하여 다를 때 NO를 출력할 수 있도록 한다.

- find_set()을 통해 현재 자신이 속한 도시를 찾고 이전 도시와 비교를 한다. 

- 다를 경우에는 NO를 answer에 대입하고 탐색을 종료한다.

answer = "YES"
before_city = -1
for city in trip_route:
    cur_city = find_set(city-1)
    if before_city == -1:
        before_city = cur_city
        continue
    if before_city != cur_city:
        answer = "NO"
        break
    before_city = cur_city

---

전체 코드

import sys
input = sys.stdin.readline

def find_set(a):
    if city_list[a] == a:
        return a
    else:
        b = find_set(city_list[a])
        city_list[a] = b
        return b

def union(a, b):
    a = find_set(a)
    b = find_set(b)
    if a != b:
        city_list[b] = a

N = int(input())
M = int(input())
graph = [list(map(int, input().split())) for _ in range(N)]
trip_route = list(map(int, input().split()))

city_list = [i for i in range(N)]
for i in range(N):
    for j in range(N):
        if graph[i][j]:
            union(i, j)

answer = "YES"
before_city = -1
for city in trip_route:
    cur_city = find_set(city-1)
    if before_city == -1:
        before_city = cur_city
        continue
    if before_city != cur_city:
        answer = "NO"
        break
    before_city = cur_city

print(answer)