백준 - 특정한 최단 경로 (1504)

문제 링크 ▶︎ 백준 특정한 최단 경로 1504

문제 전략

이 문제는 다익스트라 알고리즘으로 푼 문제이다.

노드는 양방향으로 딕셔너리에 담았다.
다익스트라 함수는 start 에서 출발하여 n까지 최단경로로 가는 visit을 리턴한다.
힙 배열을 통해 현재까지의 경로 + 다음 노드까지의 거리 < 다음 노드의 cost 라면 visit 배열을 변경해준다.

이 문제는 v1, v2를 거쳐서 n에 도달해야하므로
① 1 -> v1 -> v2 -> n
② 1 -> v2 -> v1 -> n
2가지를 비교해서 min 값을 선택한다.

코드

import sys
from heapq import heapify, heappop, heappush
input = sys.stdin.readline

def dijkstra(start):
    visit = [987654321] * (n+1)
    h = []
    heappush(h,[0,start])
    visit[start] = 0

    while h:
        dist, num = heappop(h)
        if dist > visit[num]:
            continue

        for i,j in node[num]:
            cost = dist + j
            if visit[i] > cost:
                visit[i] = cost
                heappush(h,[cost,i])

    return visit

n, e = map(int,input().split())
node = {i:[] for i in range(1,n+1)}
for j in range(e):
    a,b,v = map(int,input().split())
    node[a].append([b,v])
    node[b].append([a,v])

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

now = dijkstra(1) # 1 ~ n
next1 = dijkstra(v1) # v1 ~ n
next2 = dijkstra(v2) # v2 ~ n

answer = min(now[v1] + next1[v2] + next2[n], now[v2] + next2[v1] + next1[n])

if answer >= 987654321:
    print(-1)
else:
    print(answer)

개선 사항

다익스트라도 어려운데 v1, v2 거쳐가는 최단경로 2개 비교는 생각하기 어려움.