[Ch. 9 - 최단 경로] 미래 도시

My code 1 : Dijkstra 알고리즘 이용

import heapq
INF = int(1e9)

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

graph = [[] for i in range(n + 1)]
distance = [INF] * (n + 1)

for i in range(m):
    a, b = map(int,input().split())
    graph[a].append((b, 1))
    graph[b].append((a, 1))

x, k = map(int,input().split())
start1 = 1
start2 = k

def dijkstra(start):
    q = []
    distance[start] = 0
    heapq.heappush(q, (0, start))
    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue
        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

dijkstra(start1)
dist1 = distance[k]
dijkstra(start2)
dist2 = distance[x]

if dist1 == INF or dist2 == INF:
    print(-1)
else:
    print(dist1 + dist2)

My code 2 : Floyd-Warshall 알고리즘 이용

INF = int(1e9)

n,m = map(int,input().split())
graph = [[INF] * (n + 1) for _ in range(n + 1)]

for i in range(m):
    a, b = map(int,input().split())
    graph[a][b] = 1
    graph[b][a] = 1

x, k = map(int,input().split())

for a in range(1, n + 1):
    for b in range(1, n + 1):
        if a == b:
            graph[a][b] = 0

for k in range(1, n + 1):
    for a in range(1, n + 1):
        for b in range(1, n + 1):
            graph[a][b] = min(graph[a][b], graph[a][k] + graph[k][b])

dist1 = graph[1][k] 
dist2 = graph[k][x]

if dist1 == INF or dist2 == INF:
    print(-1)
else:
    print(dist1 + dist2)

Answer : Floyd-Warshall 알고리즘 이용

INF = int(1e9)

n,m = map(int,input().split())
graph = [[INF] * (n + 1) for _ in range(n + 1)]

for a in range(1, n + 1):
    for b in range(1, n + 1):
        if a == b:
            graph[a][b] = 0

for i in range(m):
    a, b = map(int,input().split())
    graph[a][b] = 1
    graph[b][a] = 1

x, k = map(int,input().split())

for k in range(1, n + 1):
    for a in range(1, n + 1):
        for b in range(1, n + 1):
            graph[a][b] = min(graph[a][b], graph[a][k] + graph[k][b])

distance = graph[1][k] + graph[k][x]

if distance >= INF:
    print(-1)
else:
    print(distance)

Code Review

이번 문제는 방문 판매원 A 가 1번 회사에 있을 때 k 번 회사에 방문한 뒤 x 번 회사에 방문하는 가장 빠른 최소시간을 구해보는 것이었습니다. 이는 최단 경로 문제에 해당하므로 Dijkstra 혹은 Floyd-Warshall 알고리즘을 이용하여 해결 할 수 있는데, 저는 두 가지 알고리즘 모두를 이용해 해답 코드를 작성해보았습니다.
먼저 Dijkstra 알고리즘을 이용한 Code 1 은 알고리즘을 이용하여 1번 회사에서 k 번 회사까지의 최단경로를 구해주고, k 번 회사에서 x 번 회사까지의 최단경로를 구해 두개를 더해주었습니다. 
다음으로 Floyd-Warshall 알고리즘을 이용한 Code 2 는 알고리즘을 이용하여 모든 회사에서 나머지 회사들까지의 최단 경로를 구해 2차원 리스트에 저장해주었습니다. 이렇게 하면 1번 에서 k 번, k 번에서 x 번 까지의 최단경로가 이미 리스트에 저장되어 있으므로 단순히 해당 번째 리스트 안의 값을 꺼내어 저장해주면 됩니다. 따라서 리스트에서 1에서 k, k에서 x 번 까지의 최단경로를 각각 저장하여 두개를 더해 출력해주었습니다.
Answer code 는 Code 2 와 같이 Floyd-Warshall 알고리즘을 이용하여 작성되었습니다.