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

문제

---

풀이

 

출발 지점 s와 도착 지점 e에 대해 v1, v2를 모두 지나는 최단 경로는 ①s -> v1 -> v2 -> e 또는 ②s -> v2 -> v1 -> e이다.

출발 지점을 인자로 전달받는 dijkstra() 함수를 작성해 다익스트라 알고리즘을 사용하면

def dijkstra(start):
    dist = [sys.maxsize for _ in range(N + 1)]
    visited = [False for _ in range(N + 1)]
    pq = PriorityQueue()

    dist[start] = 0
    pq.put((0, start))  # pq = (거리, 노드)

    while pq.qsize() > 0:
        d, s = pq.get()

        if visited[s]:
            continue
        visited[s] = True

        for e, nd in graph[s]:
            if dist[e] > dist[s] + nd:
                dist[e] = dist[s] + nd
                pq.put((dist[e], e))

    return dist

출발 지점부터 v1(v2)까지, v1(v2)부터 v2(v1)까지, v2(v1)부터 도착 지점까지의 최단 경로를 모두 구할 수 있으며 ①과 ② 중 더 작은 값이 정답이 된다.

 

전체 코드

import sys
input = sys.stdin.readline
from queue import PriorityQueue

N, E = map(int, input().split())
graph = [[] for _ in range(N + 1)]
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())

def dijkstra(start):
    dist = [sys.maxsize for _ in range(N + 1)]
    visited = [False for _ in range(N + 1)]
    pq = PriorityQueue()

    dist[start] = 0
    pq.put((0, start)) 

    while pq.qsize() > 0:
        d, s = pq.get()

        if visited[s]:
            continue
        visited[s] = True

        for e, nd in graph[s]:
            if dist[e] > dist[s] + nd:
                dist[e] = dist[s] + nd
                pq.put((dist[e], e))

    return dist

dist = dijkstra(1)      # 1 -> 모든 노드
from_v1 = dijkstra(v1)    # v1 -> 모든 노드
from_v2 = dijkstra(v2)    # v2 -> 모든 노드

res = min(dist[v1] + from_v1[v2] + from_v2[N], dist[v2] + from_v2[v1] + from_v1[N])
print(res) if res < sys.maxsize else print(-1)