Python) BOJ 18352. 특정 거리의 도시 찾기
문제 링크
제한
- 시간 제한 : 2 초
- 메모리 제한 : 256 MB
문제
-
어떤 나라에는 1번부터 N번까지의 도시와 M개의 단방향 도로가 존재한다. 모든 도로의 거리는 1이다.
-
이 때 특정한 도시 X로부터 출발하여 도달할 수 있는 모든 도시 중에서, 최단 거리가 정확히 K인 모든 도시들의 번호를 출력하는 프로그램을 작성하시오. 또한 출발 도시 X에서 출발 도시 X로 가는 최단 거리는 항상 0이라고 가정한다.
-
예를 들어 N=4, K=2, X=1일 때 다음과 같이 그래프가 구성되어 있다고 가정하자.
- 이 때 1번 도시에서 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 2인 도시는 4번 도시 뿐이다. 2번과 3번 도시의 경우, 최단 거리가 1이기 때문에 출력하지 않는다.
입력
- 첫째 줄에 도시의 개수 N, 도로의 개수 M, 거리 정보 K, 출발 도시의 번호 X가 주어진다. (2 ≤ N ≤ 300,000, 1 ≤ M ≤ 1,000,000, 1 ≤ K ≤ 300,000, 1 ≤ X ≤ N) 둘째 줄부터 M개의 줄에 걸쳐서 두 개의 자연수 A, B가 공백을 기준으로 구분되어 주어진다. 이는 A번 도시에서 B번 도시로 이동하는 단방향 도로가 존재한다는 의미다. (1 ≤ A, B ≤ N) 단, A와 B는 서로 다른 자연수이다.
출력
-
X로부터 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 K인 모든 도시의 번호를 한 줄에 하나씩 오름차순으로 출력한다.
-
이 때 도달할 수 있는 도시 중에서, 최단 거리가 K인 도시가 하나도 존재하지 않으면 -1을 출력한다.
👀 풀이
- 다익스트라 알고리즘을 이용해 출발점에서 각 도시까지의 최단거리를 구한 뒤 최단거리가 k와 같은 도시를 차례대로 출력했다.
전체 코드
import heapq
import sys
input = sys.stdin.readline
INF = 1e9
n, m, k, x = map(int, input().split())
graph = [[] for i in range(n+1)]
for _ in range(m):
a, b = map(int, input().split())
graph[a].append(b)
# 최단거리를 저장할 테이블
# 처음에는 모든 도시의 거리를 무한대로 초기화한다.
distance = [INF] * (n+1)
def dijkstra(start):
q = []
heapq.heappush(q, (0, start))
distance[start] = 0
while q:
dist, now = heapq.heappop(q)
if distance[now] < dist:
continue
for i in graph[now]:
cost = dist + 1
if cost < distance[i]:
distance[i] = cost
heapq.heappush(q, (cost, i))
dijkstra(x)
ans = []
for i in range(1, n+1):
if distance[i] == k:
heapq.heappush(ans, i)
if not ans:
print(-1)
else:
for node in ans:
print(node)
코드 분석
import heapq
def dijkstra(start):
q = []
heapq.heappush(q, (0, start)) # (거리, 정점)
distance[start] = 0
- 우선순위 큐를 사용해 현위치에서 최단거리가 가장 짧은 노드를 찾을 것이라서
heapq를 임포트한다. - 우선순위 큐에는 거리를 기준으로 오름차순 정렬을 할 수 있게 거리가 맨 앞으로 오는 튜플을 만들어 저장한다.
- 시작점과의 거리는 0으로 저장한 후 최단거리를 기록하는 테이블의 시작점 인덱스에도 0을 기록해준다.
while q:
dist, now = heapq.heappop(q)
# 이미 처리된 정점이면 pass
if distance[now] < dist:
continue
# 현재 정점과 연결된 정점들을 탐색한다.
for i in graph[now]:
cost = dist + 1
if cost < distance[i]:
distance[i] = cost
heapq.heappush(q, (cost, i))
- 큐의 원소를 하나씩 꺼내서 최단거리를 계산하는 작업을 시작한다.
- 만약 최단거리 테이블에 저장된 현재 노드까지의 거리가 방금 큐에서 꺼낸 거리보다 짧으면 이미 최단거리를 찾은 것이니까 더 확인해 볼 필요가 없다.
- 그렇지 않고 최단거리를 갱신해주어야 할 필요가 있는 정점이라면 현재 정점과 연결된 정점까지의 거리를 탐색한다.
- 이 문제에서 모든 정점간의 거리는 1이니까 현재 정점에서 다음 정점까지의 거리를 구하려면 시작점에서 현재 정점까지의 거리
dist에 1을 더해주면 된다. - 만약 새로 구한 거리가 최단거리 테이블에 있는 값보다 적으면 갱신한 후 우선순위 큐에 삽입한다.
- 이 과정을 큐가 빌 때까지 반복한다.
dijkstra(x)
ans = []
for i in range(1, n+1):
if distance[i] == k:
heapq.heappush(ans, i)
- 다익스트라 메서드를 호출해서
x에서부터 모든 도시까지의 최단거리를 구한다. - 우선순위 큐
ans를 만들어 최단거리가k와 같은 정점 목록을 저장한다. - 문제에서 도시 번호가 작은 순서대로 출력하라고 했기 때문에 최소힙으로 만들었다.
if not ans:
print(-1)
else:
for node in ans:
print(node)
- 만약
ans가 비어 있으면 최단거리가k인 도시가 존재하지 않는 것이니까-1을 리턴한다. - 그렇지 않으면
ans를 순회하면서 한 줄에 하나씩 출력하면 된다.
결과
- 시간 : 2784 ms
- 메모리 : 107284 KB