[백준1504/골드4] 특정한 최단 경로 - Python

https://www.acmicpc.net/problem/1504

 

처음엔 함수 없이 1 -> N 까지 다익스트라로 최단경로를 찾고, 그 과정에서 경유지 두개가 포함되는 경우에만  result 갱신하는 방법으로 풀려고 했다.

근데 경유지 안지나가고 최단경로를 찾았을 때 다른 경로를 찾지 않아서 코드 완성 못하고 끝냈다.

그래서 경유지를 사이에 끼워서 1 -> 경유지1 -> 경유지2 -> N / 1 -> 경유지2 -> 경유지1 -> N 이 두가지 경우를 따로따로 구하고 더 작은 값을 출력하도록 했다.

 

# 성공 코드

from heapq import heappush, heappop

N, E = map(int, input().split()) # 정점 수, 간선 수
adjl = [[] for _ in range(N+1)]

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

stop1, stop2 = map(int, input().split())

# start 부터 end까지의 최단거리
def dijkstra(start, end):

    heap = [(0, start)]
    visited = [float('inf')] * (N+1)
    visited[start] = 0

    while heap:

        t, c = heappop(heap)

        if c == end:
            return t

        if visited[c] < t:
            continue

        for nc, nt in adjl[c]:
            if nt + t < visited[nc]:
                visited[nc] = nt + t
                heappush(heap, (nt+t, nc))

    return float('inf')

# 1 -> stop1 -> stop2 -> N
# 1 -> stop2 -> stop1 -> N 중 더 빠른거리 채택
case1 = dijkstra(1, stop1) + dijkstra(stop1, stop2) + dijkstra(stop2, N)
case2 = dijkstra(1, stop2) + dijkstra(stop2, stop1) + dijkstra(stop1, N)

result = min(case1, case2)

print(result if result != float('inf') else -1)

 

이렇게 풀긴 풀었는데 시간이 엄청 오래걸렸다. (python 4124 ms, 66164 KB)

다익스트라를 6번이나 구해서 그런 거 같은데 최적화 고민을 해봐야할 거 같다.