[백준/Python] 특정 거리의 도시 찾기

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

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

처음에 제출했을 때 시간 초과가 났다! 이유를 알아보니 파이썬은 input().split() 으로 입력하는 경우 입력량이 많기 때문에 시간 초과가 나는 것이었다... 시간 초과가 나지 않으려면 sys.stdin.readline() 을 사용해야 한다!

 

 

이 문제는 최단 거리를 찾는 문제로 너비 우선 탐색(BFS)을 이용하면 풀 수 있었다!

import sys
from collections import deque

input = sys.stdin.readline
# 입력 받는 부분 
n, m, k, x = map(int, input().split()) # 도시의 개수, 도로의 개수, 거리 정보, 출발 도시 번호

graph = [[] for _ in range(n+1)] # 연결된 도로 정보 리스트
for i in range(m):
	a, b = map(int, input().split())
	graph[a].append(b)

distance = [-1] * (n+1) # 도시에 대한 최단거리 리스트
queue = deque([x]) # 너비 우선 탐색을 위한 큐
distance[x] = 0 # 출발 도시의 거리는 0으로

# 너비 우선 탐색 수행
while queue:
	v = queue.popleft() # 현재 도시

    # 현재 도시에서 이동할 수 있는 모든 도시들을 확인
	for i in graph[v]:
		if distance[i] == -1: # 방문하지 않은 도시라면
			queue.append(i)
			distance[i] = distance[v] + 1 # 최단 거리 갱신

# 최단 거리가 k인 도시 번호 출력
city = False
for i in range(1, n+1):
	if distance[i] == k:
		print(i)
		city = True

# 최단 거리가 k인 도시가 없다면 -1 출력
if city == False:
	print(-1)

도시 번호가 1부터 시작하므로 도시번호와 리스트의 인덱스를 맞추기 위해 리스트는 n+1 까지 생성해준다. 너비 우선 탐색을 할 때는 deque를 사용한다. queue 가 없을 때까지 너비 우선 탐색을 수행한다. 현재 도시 정보를 queue 에서 가져와 현재 도시에서 이동할 수 있는 모든 도시들을 확인한다. 방문하지 않은 도시라면 최단 거리를 갱신한다. 너비 우선 탐색이 끝나면 최단 거리가 k인 도시번호를 출력하면 된다.