[백준/파이썬] 1753번: 최단경로

문제

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

풀이

• 방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로 출력
• 시작점 K에서 모든 정점까지의 최단 거리
  • 도달할 수 없으면 INF 출력

다익스트라 알고리즘을 이용해 푸는 최단 경로 문제입니다.
heapq를 사용해 문제를 풀어보겠습니다.

if __name__ == "__main__": # 메인문
    a,b = map(int,input().split())
    k = int(input())
    graph = [[] for _ in range(a+1)]
    dist = [1e9] * (a+1)
    
    for _ in range(b):
        u,v,w = map(int,input().split()) # u -> v, 가중치 w
        graph[u].append((v,w))

• 정점의 개수 a, 간선의 개수 b, 시작점 k
• b개 줄에 걸쳐 그래프의 정보 입력

def func(start):
    queue = [] # heapq 선언
    dist[start] = 0 # 자신으로부터 거리 0
    heapq.heappush(queue, (0, start)) # 거리, 노드
    
    while queue:
        dis, cur = heapq.heappop(queue)
        
        if dis > dist[cur]: # 더 짧은 경로가 존재하면 pass
            continue
        
        for node, weight in graph[cur]:
            w = weight + dis # 현재까지 거리 + 가중치
            if w < dist[node]: # w가 다음 노드 경로보다 작을 경우 업데이트
                dist[node] = w
                heapq.heappush(queue, (w, node))

• 다익스트라 함수

코드

import sys, heapq
input = sys.stdin.readline

def func(start):
    queue = []
    dist[start] = 0
    heapq.heappush(queue, (0, start))
    
    while queue:
        dis, cur = heapq.heappop(queue)
        
        if dis > dist[cur]:
            continue
        
        for node, weight in graph[cur]:
            w = weight + dis
            if w < dist[node]:
                dist[node] = w
                heapq.heappush(queue, (w, node))

if __name__ == "__main__":
    a,b = map(int,input().split())
    k = int(input())
    graph = [[] for _ in range(a+1)]
    dist = [1e9] * (a+1)
    
    for _ in range(b):
        u,v,w = map(int,input().split())
        graph[u].append((v,w))
    
    func(k)
    
    for i in range(1, a+1):
        print(dist[i] if dist[i] != 1e9 else "INF")