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

문제

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

입력

• 첫째 줄에 도시의 개수 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을 출력한다.

문제 풀이

• 노드 1에서 출발해서 도달할 수 있는 모든 거리의 정보를 탐색하기 위하여 다익스트라 알고리즘 사용
• 다익스트라 알고리즘을 사용하여 1부터 모든 노드들까지의 거리를 distance 리스트에 저장
• distance에 k값이 있다면 그 인덱스를 출력하고, 없다면 -1출력

코드

import heapq
import sys
input = sys.stdin.readline
n, m, k, x = map(int,input().split())
INF = int(1e9)
graph = [[] for _ in range(n + 1)]
distance = [INF] * (n + 1)

# 간선 정보 입력받기
for _ in range(m):
    a, b = map(int,input().split())
    graph[a].append((b, 1))


# 다익스트라 알고리즘 수행
def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        # 거리, 노드 pop
        dist, now = heapq.heappop(q)
        # 꺼낸 거리가 더 크다면 무시
        if distance[now] < dist:
            continue
        # 꺼낸 거리가 더 작다면 탐색할 곳과 연결된 노드들 탐색
        for i in graph[now]:
            # 그래프 = (노드, 거리)
            cost = dist + i[1]
            # 현재 위치한 곳을 거쳐가는 경우의 값이 더 작다면 거리 값 갱신
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

dijkstra(x)
cnt = 0
for i in range(1, n + 1):
    if distance[i] == k:
        cnt += 1
        print(i)
    elif cnt == 0 and i == n:
        print(-1)