[백준] 1504번 특정한 최단경로 - python파이썬 (시간초과 해결)

문제

 

 

 

 

 

처음 이 문제를 보았을 때  한 정점에서 다른 모든 정점까지의 거리가 아니라 1번정점에서 N번정점으로 가는 최단거리를 물어보다 보니 플로이드-워샬 알고리즘으로 풀어주려고 했다.

 

 

하지만 N의 범위가 800까지로 N의 세제곱이되는 플로이드-워샬의 시간복잡도를 생각하면 시간초과가 우려된다는 것을 알 수 있다.

그리고 역시나 시간초과가 떴다.

 

 

그래서 어떻게 하나 고민하던 찰나 이 문제에는 특정한 조건이 있다 바로 v1, v2 노드를 반드시 거쳐서 가야한다는 점

 

 

그렇게 되면 가능한 루트는 딱 두가지가 된다.

 

1.  1번노드 -> v1 -> v2 -> N
2. 1번노드 -> v2 -> v1 -> N

 

 

 

그렇기 때문에 다익스트라 알고리즘으로도 풀 수가 있다.

 

 

이게 무슨 뜻이냐면은 다익스트라 알고리즘은 결국 한 정점에서 모든정점과의 최단거리를 구하는 것이다

그렇기때문에 위의 두가지 루트의 최단거리를 구하는 데 필요한 각각의 최단거리를 전부 구해줄 수 있다는 것이다.

 

우리가 필요한 것은 (1,v1) (v1,v2) (v2,N) (1,v2) (v2,v1) (v1,N) 이 노드들간의 최단거리이다. 시작노드를 적절히 변경하면서 구해준 다음  두가지 루트중 더 짧은 것을 정답으로 지정해주면 된다.

 

 

 

정답코드

 

import heapq

N, E = map(int,input().split())
inf = float("INF")
graph = [[] for i 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,end):
    distance = [inf] * (N+1)
    q = []
    heapq.heappush(q,(start,0))
    distance[start] = 0

    while q:
        now, dist  = heapq.heappop(q)
        #이미 처리된 노드는 무시
        if distance[now] < dist:
            continue
        #현재 노드와 인접한 노드 확인
        for next_node,ndist in graph[now]:
            d = dist + ndist
            if d < distance[next_node]:
                distance[next_node] = d
                heapq.heappush(q,(next_node,d))
    
    return distance[end]


ans = min(dijkstra(1,v1)+dijkstra(v1,v2)+dijkstra(v2,N),dijkstra(1,v2)+dijkstra(v2,v1)+dijkstra(v1,N))

if ans < inf:
    print(ans)
else:
    print(-1)

 

다익스트라와 플로이드-워샬 알고리즘을 공식 외우듯이 외운 사람은 더 생각하는게 어려웠을 법한 문제이다. 너무 알고있는 것에 갇혀서 생각하면 안된다는 것을 깨닫게 해준 문제였다.