[BOJ] 1504. 특정한 최단 경로(python)

📌문제

💪아이디어

1. 양방향 그래프
   defaultdict를 사용해서 노드 간 간선을 이어준다.
   

2. 1번에서 n번 최단거리
   heap을 사용한 다익스트라알고리즘을 이용한다. 1에서 n까지 가는 최단 거리에 두 정점을 반드시 통과해야한다.
   
   • 1) start -> v1 -> v2 -> N
     dist[start][v1], dist[v1][v2],dist[v2][end]
     
   • 2) start -> v2 -> v1 -> end
     dist[start][v2], dist[v2][v1],dist[v1][end]
     
     

   양방향 그래프이기 때문에 v1 -> v2와 v2 -> v1는 같다. 따라서, start를 시작점으로하는 다익스트라와 v1를 시작점으로하는 다익스트라, v2를 시작점으로하는 다익스트라 총 3개의 다익스트라 배열을 구하면 된다.
   
   

3. 두 개의 정점을 지나는 최단 경로의 길이
   위에서 나눈 두가지 경우의 길이를 모두 더해서 최솟값을 출력한다. 최솟값이 inf라면 -1를 출력한다.
   

🥂코드

import sys; input = sys.stdin.readline
from collections import defaultdict
from heapq import heappush, heappop

def Dijk(start):
    dist = [float('inf')] * (n+1)
    dist[start] = 0
    heap = []
    heappush(heap,(0,start))

    while heap:
        d, cur = heappop(heap)
        if dist[cur] < d: continue
        for nxt, cost in graph[cur]:
            distance = cost + d
            if distance < dist[nxt]:
                dist[nxt] = distance
                heappush(heap,(distance,nxt))
    return dist

n, e = map(int,input().split())
graph = defaultdict(list)

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())

startDijk = Dijk(1)
start2v1 = startDijk[v1]
start2v2 = startDijk[v2]

nodeDijk = Dijk(v1)
v12v2 = nodeDijk[v2]
v12end = nodeDijk[n]

v2Dijk = Dijk(v2)
v22end = v2Dijk[n]

answer = min((start2v1 + v12v2 + v22end), (start2v2 + v12v2 + v12end))
print(answer if answer != float('inf') else -1)