백준 알고리즘 1504번 특정한 최단 경로(python)

기본적인 다익스트라 문제이다.

다익스트라를 1(맨 처음 노드), v1, v2 기준으로 3번 돌려서 3개의 최단 거리 테이블을 생성하고

1 -> v1 -> v2 -> n 순서로 갈 때와 1 -> v2 -> v1 -> n 순서로 갈 때의 거리 중 더 짧은 길이를 출력해주면 된다.

import heapq
import sys
INF = sys.maxsize
input = sys.stdin.readline

def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    dis = [INF]*(N+1)
    dis[start] = 0
    while q:
        d, now = heapq.heappop(q)
        if dis[now] < d:
            continue
        for v, w in graph[now]:
            cost = d + w
            if cost < dis[v]:
                dis[v] = cost
                heapq.heappush(q, (cost, v))
    return dis

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))
v1, v2 = map(int, input().split())
d1 = dijkstra(1)
d2 = dijkstra(v1)
d3 = dijkstra(v2)
res = min(d1[v1]+d2[v2]+d3[N], d1[v2]+d3[v1]+d2[N])
print(res if res < INF else -1)