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

문제

방향성이 없는 그래프가 주어진다. 세준이는 1번 정점에서 N번 정점으로 최단 거리로 이동하려고 한다. 또한 세준이는 두 가지 조건을 만족하면서 이동하는 특정한 최단 경로를 구하고 싶은데, 그것은 바로 임의로 주어진 두 정점은 반드시 통과해야 한다는 것이다.

세준이는 한번 이동했던 정점은 물론, 한번 이동했던 간선도 다시 이동할 수 있다. 하지만 반드시 최단 경로로 이동해야 한다는 사실에 주의하라. 1번 정점에서 N번 정점으로 이동할 때, 주어진 두 정점을 반드시 거치면서 최단 경로로 이동하는 프로그램을 작성하시오.

입력

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

출력

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

예제 입력 1

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

예제 출력 1

7

나의 풀이

앞서 1753번 최단경로 문제에서 다익스트라 알고리즘을 사용해 특정 노드에서 다른 노드들로 가는 최단경로들을 구하는 문제를 풀었었다. 이 문제는 다익스트라 알고리즘을 사용해서 각 정점들의 최단경로들을 구한 뒤 경우의 수를 따져 최솟값을 구해주면 된다.

1번 노드에서 출발하여 N까지 가는데 거쳐야하는 두 노드를 v1, v2로 놓아보자. 그럼 이 두 노드를 거쳐 N까지 가는 경로는 두 가지의 경우의 수가 주어진다. 1 -> v1 -> v2 -> N과 1 -> v2 -> v1 -> N이다. 따라서 이 두 가지 경우의 경로에서 최솟값들을 구해 이 중에서 작은 값이 곧 답이 된다.

그럼 이 두 가지의 최솟값은 어떻게 구하면 될까? 두 가지 중 한 가지인 1 -> v1 -> v2 -> N을 예로들자면 1번에서 v1으로 가는 최솟값과 v1에서 v2로 가는 최솟값, v2에서 N으로 가는 최솟값을 더해주면 되지 않을까? 그럼 이 최솟값은 자연스럽게 다익스트라 알고리즘을 사용해서 구하면 된다는 생각을 할 수 있다. 따라서 1번에서 출발하는 다익스트라 알고리즘. v1에서 출발하는 다익스트라 알고리즘, v2에서 출발하는 다익스트라 알고리즘을 돌려 각 노드로 이동하는 최솟값을 담는 배열을 만들어놓고 위에서 말한 두 가지의 경우의 값을 계산하면 된다.

이 문제에서는 무방향그래프가 주어지므로 인접리스트를 만들때 1 2 3이 주어진다면 1번과 2번 노드에 각각 [2,3], [1,3]을 넣어주어야 한다는 것도 주의해야한다.

코드

# 1504번 특정한 최단 경로
import sys, heapq

inf = float('inf')

# dijkstra
def dijkstra(start, g, n):
    # 최소 비용 배열
    d = [inf] * (n+1)
    d[start] = 0

    # 최소 힙
    mheap = []
    heapq.heappush(mheap, [0,start])

    while mheap:
        tmp = heapq.heappop(mheap)
        cost = tmp[0]
        node = tmp[1]

        if cost > d[node]:
            continue

        for n in g[node]:
            neighbor = n[0]
            n_cost = d[node] + n[1]

            if  n_cost < d[neighbor]:
                d[neighbor] = n_cost
                heapq.heappush(mheap, [n_cost,neighbor])
    return d        

# mainㅌ
n, e = [int(x) for x in sys.stdin.readline().split()]

# 인접리스트
g = [[] for _ in range(n+1)]
for _ in range(e):
    a, b, c = [int(x) for x in sys.stdin.readline().split()]
    # 무방향 그래프로 양쪽 모두 추가
    g[a].append([b,c])
    g[b].append([a,c])

v1, v2 = [int(x) for x in sys.stdin.readline().split()]

d_start = dijkstra(1,g,n)
d_v1 = dijkstra(v1,g,n)
d_v2 = dijkstra(v2,g,n)

answer = min(d_start[v1] + d_v1[v2] + d_v2[n], d_start[v2] + d_v2[v1] + d_v1[n])
print(answer if answer != inf else -1)