[백준 1504] 특정한 최단 경로

문제 

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

 

코드

import sys, heapq, math

def Input_Data():
    input = sys.stdin.readline
    N, E = map(int, input().split())
    graph = {i : [] for i in range(1, N+1)}
    for _ in range(E):
        # 양방향 그래프 관계 저장
        a,b,c = map(int, input().split())
        graph[a].append([b,c]) 
        graph[b].append([a,c])

    A, B = map(int, input().split())
    return N, E, graph, A, B

def Dijkstra(start, end):
    dist = [math.inf] * (N+1)
    heap = []
    heapq.heappush(heap,(0,start)) 
    dist[start] = 0 
    while heap: 
        cst, cur = heapq.heappop(heap)
        if cst > dist[cur]: #현재 노드보다 먼 거리는 볼 필요 없음
            continue 
        for next, cost in graph[cur]:
            if cst + cost < dist[next]: #앞으로 갈 곳의 비용이 (현재 + 다음 노드) 가 더 작으면 
                dist[next] = cst + cost
                heapq.heappush(heap,(cst + cost, next))
    return dist[end]
        
N, E, graph, A, B = Input_Data()
# 1 -> A -> B -> N
start = Dijkstra(1,A)
via = Dijkstra(A,B)
end = Dijkstra(B,N)

res1 = start + via + end
# 1 -> B -> A -> N
start = Dijkstra(1,B)
via = Dijkstra(B,A)
end = Dijkstra(A,N)

res2 = start + via + end

# 도달 할 수 없는 경우 
if res1 >= math.inf and res2>= math.inf:
    print(-1)
else:
    print(min(res1, res2))