백준 1504 특정한 최단 경로 (파이썬)

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

1504번: 특정한 최단 경로

거쳐야하는 두 개의 정점을 a, b라고 하면

시작 - a - b - 끝, 시작 - b - a -끝 두 케이스에 대한 경로를 다익스트라로 구하면 된다

그리고 발생할 수 있는 모든 케이스를 고려하여 출력이 제대로 나올 수 있도록 설계

from heapq import heappush, heappop
import sys

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

def dij(x):
    d = [INF]*n
    heappush(q, [0, x])
    d[x] = 0
    while q:
        w, x = heappop(q)
        for nw, nx in a[x]:
            nw += w
            if nw < d[nx]:
                d[nx] = nw
                heappush(q, [nw, nx])
    return d

n, e = map(int, input().split())
a = [[] for _ in range(n)]
q = []

for _ in range(e):
    x, y, w = map(int, input().split())
    a[x-1].append([w, y-1])
    a[y-1].append([w, x-1])

x, y = map(int, input().split())
d0 = dij(0); d1 = dij(x-1); d2 = dij(y-1)
ans1, ans2, flag1, flag2 = sys.maxsize, sys.maxsize, 0, 0
if d0[x-1] != INF and d1[y-1] != INF and d2[n-1] != INF:
    ans1 = d0[x-1] + d1[y-1] + d2[n-1]
else:
    flag1 = 1
if d0[y-1] != INF and d2[x-1] != INF and d1[n-1] != INF:
    ans2 = d0[y-1] + d2[x-1] + d1[n-1]
else:
    flag2 = 1
if flag1 and flag2:
    print(-1)
else:
    print(min(ans1, ans2))