[골드 4] 백준 1976 - 여행 가자 (파이썬)

https://www.acmicpc.net/problem/1976

풀이

1. N개의 도시가 있고, 어떤 도시끼리는 도로로 연결되어 있다.
2. M개의 도시가 주어진 여행 계획에 포함되어 있다.
3. 주어진 도시들이 모두 같은 연결 요소에 속해 있다면 여행이 가능하다.
4. 연결 정보는 인접 행렬로 주어진다.

이 문제는 그래프 연결 요소 판별 문제이다.

• BFS를 통해 각 도시가 어떤 연결 요소에 속해 있는지 visit 배열에 저장하고
• 여행 계획에 포함된 도시들이 모두 같은 연결 요소에 속해 있는지를 확인한다.
• 하나라도 다르면 NO, 모두 같으면 YES를 출력한다.

N = int(input())   # 도시 수
M = int(input())   # 여행에 포함된 도시 수

adj = [[] for _ in range(N)]
for i in range(N) :
    A = list(map(int, input().split()))
    for j in range(N) :
        if A[j] == 1 :       # 연결되어 있다면
            adj[i].append(j)

인접 행렬을 인접 리스트로 변환하고, adj[i]는 도시 i와 연결되 도시 목록을 나타낸다.

T = list(map(int, input().split()))
for i in range(len(T)) :
    T[i] -= 1   # 0-indexed로 바꿈

이해하기 쉽도록 여행 경로를 0-based 인덱스로 변경한다.

for i in range(N) :
    if visit[i] == -1 :
        queue.append(i)
        visit[i] = i

        while len(queue) != 0 :
            u = queue.popleft()
            for v in adj[u] :
                if visit[v] == -1 :
                    queue.append(v)
                    visit[v] = i

BFS로 모든 연결 요소에 번호를 부여하고, 같은 번호면 같은 그룹에 있는 것이다.

possible = True
for i in range(1, len(T)) :
    if visit[T[i - 1]] != visit[T[i]] :
        possible = False
        break

여행 계획 도시들이 모두 같은 연결 요소에 있는지 확인하고, 답을 출력한다.

코드

# 백준 1976 - 여행 가자
# 분류 : BFS

from collections import deque

N = int(input())
M = int(input())

adj = [[] for _ in range(N)]
for i in range(N) :
    A = list(map(int, input().split()))

    for j in range(N) :
        if A[j] == 1 :
            adj[i].append(j)

T = list(map(int, input().split()))
for i in range(len(T)) :
    T[i] -= 1

visit = [-1] * N
queue = deque()

for i in range(N) :
    if visit[i] == -1 :
        queue.append(i)
        visit[i] = i

        while len(queue) != 0 :
            u = queue.popleft()

            for v in adj[u] :
                if visit[v] == -1 :
                    queue.append(v)
                    visit[v] = i

possible = True
for i in range(1, len(T)) :
    if visit[T[i - 1]] != visit[T[i]] :
        possible = False
        break

if possible :
    print("YES")
else :
    print("NO")