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

문제

1504번: 특정한 최단 경로

입력

• 첫째 줄에 정점의 개수 N과 간선의 개수 E가 주어진다. (2 ≤ N ≤ 800, 0 ≤ E ≤ 200,000) 둘째 줄부터 E개의 줄에 걸쳐서 세 개의 정수 a, b, c가 주어지는데, a번 정점에서 b번 정점까지 양방향 길이 존재하며, 그 거리가 c라는 뜻이다. (1 ≤ c ≤ 1,000) 다음 줄에는 반드시 거쳐야 하는 두 개의 서로 다른 정점 번호 v1과 v2가 주어진다. (v1 ≠ v2, v1 ≠ N, v2 ≠ 1) 임의의 두 정점 u와 v사이에는 간선이 최대 1개 존재한다.

출력

• 첫째 줄에 두 개의 정점을 지나는 최단 경로의 길이를 출력한다. 그러한 경로가 없을 때에는 -1을 출력한다.

입력 예제

4 6
1 2 3
2 3 3
3 4 1
1 3 5
2 4 5
1 4 4
2 3

출력 예제

7

문제 풀이

• 다익스트라 알고리즘을 수행하여 첫번째 노드에서 출발하는 경우, v1에서 출발하는 경우, v2에서 출발하는 경우의 모든 최단거리를 구한다
• 그 후, v1과 v2를 거쳐 n까지 도달하는 최소 거리를 구한다
  • 1 -> v1 -> v2 -> n
  • 1 -> v2 -> v1 -> n
• v1과 v2를 거쳐 n에 도달하지 못하는 경우 -1을 출력하고, 도달 할 수 있다면 최단경로를 출력한다.

코드

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)
n, m = map(int,input().split())
graph = [[] for _ in range(n + 1)]

for _ in range(m):
    a,b,c = map(int,input().split())
    graph[a].append((b, c))
    graph[b].append((a, c))

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

def dijkstra(start):
    distance = [INF] * (n + 1)
    q = []
    heapq.heappush(q,(0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue
        for i in graph[now]:
            cost = dist + i[1]
            if distance[i[0]] > cost:
                distance[i[0]] = cost
                heapq.heappush(q,(cost, i[0]))
    return distance

start = dijkstra(1) # 출발정점이 1일 때의 모든거리
first = dijkstra(v1) # 출발정점이 2일 때의 모든거리
second = dijkstra(v2) # 출발정점이 3일 때의 모든거리

# 1 -> 2 , 2 -> 3, 3 -> 4
path1 = start[v1] + first[v2] + second[n]
# 1 -> 3, 3 -> 2, 2 -> 4
path2 = start[v2] + second[v1] + first[n]
result = min(path1, path2)

if result >= INF:
    print(-1)
else:
    print(result)