[백준1753번, 다익스트라 알고리즘] 최단 경로 - Python

타입 : 다익스트라 알고리즘

문제 : 방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

입력 : 첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

출력 : 첫째 줄부터 V개의 줄에 걸쳐, i번째 줄에 i번 정점으로의 최단 경로의 경로값을 출력한다. 시작점 자신은 0으로 출력하고, 경로가 존재하지 않는 경우에는 INF를 출력하면 된다.

import heapq # 힙 자료구조
import sys # 입력을 위한

Inf = int(1e9) # 가장 큰 값 설정

input = sys.stdin.readline # input 함수 재정의

v, e = map(int, input().split())  # vertex, edge 입력

distances = [Inf] * (v + 1) # v의 vertex의 distance를 모두 Inf로 초기화

graph = [[] for _ in range(v + 1)] # graph 선언

start = int(input())  # start point 입력

for _ in range(e): # e 개의 edge에 대해서 edge 정보 입력
    a, b, c = map(int, input().split()) 
    graph[a].append((b, c)) # graph[a]에 (b,c)를 tuple 형태로 입력


def djikstra(start): # 힙 자료구조를 이용한 다익스트라 알고리즘
    q = []
    heapq.heappush(q, (0, start))
    distances[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        if distances[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + i[1]
            if cost < distances[i[0]]:
                distances[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))


djikstra(start) # 다익스트라 알고리즘 실행

for i in range(1, v + 1):  # 0은 고려하지 않음으로 1부터 n까지 distance를 출력
    if distances[i] == Inf: # 만약 경로가 존재하지 않는 경우
        print("INF") # INF 출력
    else: # 그렇지 않은 경우,
        print(distances[i]) # distance 출력

실행 결과 :

5 6
1
5 1 1
1 2 2
1 3 3
2 3 4
2 4 5
3 4 6
0
2
3
7
INF