백준 1504번 : 특정한 최단 경로 (Python)

출처 : https://www.acmicpc.net/problem/1504

import heapq
# import sys
# input = sys.stdin.readline

n, e = map(int, input().split())
graph = [[] for _ in range(n + 1)]
for _ in range(e) :
    a, b, c = map(int, input().split())
    graph[a].append((b, c))
    graph[b].append((a, c))
# graph = [[],
#          [(2, 3), (3, 5), (4, 4)],
#          [(1, 3), (3, 3), (4, 5)],
#          [(2, 3), (4, 1), (1, 5)],
#          [(3, 1), (2, 5), (1, 4)]]

v1, v2 = map(int, input().split())

INF = int(1e9)
distance = [INF] * (n + 1)

def dijkstra(graph, start) :
    q = []
    distance = [INF] * (n + 1)
    distance[start] = 0
    heapq.heappush(q, (0, start))
    while q :
        dist_now, now = heapq.heappop(q)
        # 큐에서 꺼낸 노드까지의 최단거리가 기록된 최단거리보다 길면,
        # 이 노드를 거쳐가는 최단거리는 계산할 필요 없음
        if dist_now > distance[now] :
            continue
        # 인접 노드 탐색
        for i in graph[now] :
            dist_adjacent = dist_now + i[1]
            # 현 노드를 거쳐갈때 거리가 더 짧으면 최단거리 업데이트, 큐에 삽입
            if dist_adjacent < distance[i[0]] :
                distance[i[0]] = dist_adjacent
                heapq.heappush(q, (dist_adjacent, i[0]))
    return distance

start_dist = dijkstra(graph, 1)
v1_dist = dijkstra(graph, v1)
v2_dist = dijkstra(graph, v2)

# v1을 먼저 거쳐가는 경로와 v2를 먼저 거쳐가는 경로의 최단거리를 비교
sol1 = start_dist[v1] + v1_dist[v2] + v2_dist[n]
sol2 = start_dist[v2] + v2_dist[v1] + v1_dist[n]
sol = min(sol1, sol2)

if sol >= INF :
    print(-1)
else :
    print(sol)

노드가 최대 800개나 되기 때문에 플로이드 워셜 알고리즘을 사용하면 시간초과가 뜬다.

따라서 다익스트라 알고리즘을 반복해서 사용해 다음 노드로 가는 최단거리를 구해 더한다.