1504: 특정한 최단 경로

풀이 시간

• 30m

구현 방식

• 1번 정점에서 N번 정점으로 최단 거리로 이동할 때, v1과 v2를 반드시 통과해야하기 때문에, 두 가지 경우를 생각해볼 수 있다
  1) 1 -> v1 -> v2 -> N
  2) 1 -> v2 -> v1 -> N

• 따라서 1번에서 N까지 가는 dijkstra, v1에서 N까지 가는 dijkstra, v2에서 N까지 가는 dijkstra를 수행한 후, 최단거리를 구해주면 된다

코드

import sys
import heapq


def dijkstra(start):
    distances = [int(10e9) for _ in range(N+1)]
    distances[start] = 0
    heap = []
    heapq.heappush(heap, [distances[start], start])

    while heap:
        curr_dist, curr = heapq.heappop(heap)
        if distances[curr] < curr_dist: continue
        for next, next_dist in graph[curr]:
            dist = curr_dist + next_dist
            if distances[next] > dist:
                distances[next] = dist
                heapq.heappush(heap, [distances[next], next])
    return distances

N, E = map(int, sys.stdin.readline()[:-1].split())
graph = [[] for _ in range(N+1)]
for e in range(E):
    a, b, c = map(int, sys.stdin.readline()[:-1].split())
    graph[a].append((b, c))
    graph[b].append((a, c))
v1, v2 = map(int, sys.stdin.readline()[:-1].split())

distance0 = dijkstra(1)
distance1 = dijkstra(v1)
distance2 = dijkstra(v2)

# 1 -> v1 -> v2 -> N
v1v2 = distance0[v1] + distance1[v2] + distance2[N]
# 1 -> v2 -> v1 -> N
v2v1 = distance0[v2] + distance2[v1] + distance1[N]

result = min(v1v2, v2v1)
if result >= int(10e9): print(-1)
else: print(result)

결과