백준 특정 거리의 도시 찾기 코드 및 해설 (파이썬)

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

18352번: 특정 거리의 도시 찾기

 

큐를 이용한 BFS 코드를 활용해 풀이했습니다. 처음엔 X에서 각 도시까지의 거리를 충분히 큰 값 INF로 초기화합니다. 여기선 문제에서 주어진 K의 최댓값인 300000보다 1만큼 큰 값을 INF로 설정했습니다. 문제에서 X에서 X까지의 거리는 0이라고 했으므로 X번째 요소는 0으로 값을 변경합니다.

 

한 번 방문한 곳을 계속해서 다시 가지 않도록 visited를 만들어주고 모두 False로 초기화합니다. 단방향 도로의 정보는 딕셔너리로 저장합니다. 하나의 도시에 여러 개의 도로가 있을 수 있으므로 value는 리스트 형태로 만들었습니다.

이후 bfs 함수를 통해 탐색을 시작합니다. 시작하는 지점을 큐에 넣은 후 큐가 빌 때까지 다음을 반복합니다.

큐의 첫 번째 요소를 팝하고, 그 요소(도시)와 이어진 도로가 있으면 그 도로를 통해 이어진 도시에 방문하고, 그 도시를 아직 방문하지 않은 상태면 visited를 True로 바꿔주고, 해당 도시까지의 거리를 [팝한 요소까지의 거리 + 1]로 갱신합니다. 그리고 이 도시를 큐에 추가합니다.

 

이렇게 갱신된 거리 리스트 중에서 K값과 같은 것들을 찾고, K값과 같은 것이 없으면 -1을, 아니면 해당 도시들을 차례로 출력합니다.

 

import sys
from collections import defaultdict, deque

# N: 도시의 개수
# M: 도로의 개수
# K: 타겟 최단 거리
# X: 출발 도시
N, M, K, X = map(int, sys.stdin.readline().split())

INF = 300001
# X에서 각 도시까지의 거리 INF로 초기화
dist = [INF for _ in range(N+1)]
# X에서 X까지의 거리는 0
dist[X] = 0
# 재방문 방지
visited = [False for _ in range(N + 1)]

# 단방향 도로의 출발지가 key, 도착지가 value
roads = defaultdict(list)
for _ in range(M):
    k, v = map(int, sys.stdin.readline().split())
    roads[k].append(v)

def bfs(start):
    q = deque([start])
    visited[start] = True
    while q:
        cur = q.popleft()
        for node in roads[cur]:
            if not visited[node]:
                visited[node] = True
                dist[node] = min(dist[node], dist[cur] + 1)
                q.append(node)
    return dist

result = bfs(X)

answer = [i for i,d in enumerate(result) if d == K]
if not answer:
    print(-1)
else:
    for ans in answer:
        print(ans)