[Python] 백준 1504번 특정한 최단 경로(dijkstra)

문제

풀이

처음에는 플로이드 와샬로 풀 생각을 하였다. 거쳐가야 할 정점이 두 개나 있기 때문이다.

하지만 그렇게 된다면 쓸모 없는 정점 까지의 최단 거리도 계산해야 한다.

따라서 플로이드 와샬은 O(n3)이 된다.

다른 언어는 플로이드 와샬도 되는 것 같은데 파이썬은 안된다. 따라서 다익스트라 알고리즘을 3번 사용하여 

시작점 -> 경유지 1 한번, 경유지1-> 경유지2  한번, 경유지-> 도착 점 한번 

이렇게 3번을 사용하거나, 경유지 1,2의 순서를 바꿔서 3번 사용하여 각각의 최단 거리를 비교하여 구한다.

처음에 시도했던 플로이드 와샬 알고리즘이다.

(시간 초과가 나는 코드)

INF = int(1e9)

n, e = map(int, input().split())

table = [[INF for _ in range(n)]for _ in range(n)]

print(table)

for i in range(e):

    a, b, cost = map(int, input().split())

    table[a-1][b-1] = cost

    table[b-1][a-1] = cost

haveToA, haveToB = map(int, input().split())

for i in range(n):

    table[i][i] = 0

for i in range(n):

    for j in range(n):

        for k in range(n):

            table[j][k] = min(table[j][k], table[j][i]+table[i][k])

print(min(table[0][haveToA-1]+table[haveToA-1][haveToB-1]+table[haveToB-1]

      [n-1], table[0][haveToB-1]+table[haveToB-1][haveToA-1]+table[haveToA-1][n-1]))

이것이 정답으로 인정된 코드이다.

import sys

import heapq

input=sys.stdin.readline

INF = int(1e9)

n, e = map(int, input().split())

table = [[]for _ in range(n)]

for i in range(e):

    a, b, cost = map(int, input().split())

    table[a-1].append((b-1, cost))

    table[b-1].append((a-1, cost))

haveToA, haveToB = map(int, input().split())

def dijkstra(start):

    q = []

    distance = [INF]*n

    heapq.heappush(q, (0, start))

    distance[start] = 0

    while q:

        dist, now = heapq.heappop(q)

        if distance[now] < dist:

            continue

        for i in table[now]:

            cost = dist + i[1]

            if cost < distance[i[0]]:

                distance[i[0]] = cost

                heapq.heappush(q, (cost, i[0]))

    return distance

result = min(dijkstra(0)[haveToA-1]+dijkstra(haveToA-1)[haveToB-1]+dijkstra(haveToB-1)

             [n-1], dijkstra(0)[haveToB-1]+dijkstra(haveToB-1)[haveToA-1]+dijkstra(haveToA-1)[n-1])

print(result if result < INF else -1)

또한 sys.stdin.readline을 사용하지 않아도 시간 초과가 난다...(pypy3 는 통과한다.)

백준에서 알고리즘을 할 때에는 반드시 저것을 사용해야겠다.