최단 경로 알고리즘 예제: 미래 도시

n개의 회사가 있고 m쌍의 회사가 연결되어있다

(양방향으로 연결되어있다)

 

1번회사에서 출발해 k번 회사를 거쳐 x번 회사까지 가는데 걸리는 최단시간은?

 

---

나는 다익스트라 알고리즘을 이용해 풀었다

1에서 출발하여 k까지 가는 최단거리 + k에서 출발하여 x까지 가는 최단거리 라고 생각할 수 있고

다익스트라를 두 번 실행하면 답을 얻을 수 있다

import sys, heapq
input=sys.stdin.readline

n,m=map(int,input().split()) #노드의 개수n, 간선개수m

graph=[[]for i in range(n+1)]
for i in range(m):
    a,b=map(int,input().split()) #서로 연결된 두 회사 a,b
    graph[a].append(b)
    graph[b].append(a)

x,k=map(int,input().split()) #목적지x 경유지k

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

dijkstra(1)
to_k=distance[k]

distance=[INF]*(n+1)
dijkstra(k)
k_to_x=(distance[x])

if to_k==INF or k_to_x==INF:
    print(-1)
else:
    print(to_k + k_to_x)

dijkstra(1)을 한 후 distance[k]를 받고,

distance리스트를 초기화 한 후 dijkstra(k)를 수행하고 distance[x]를 받는다

 

근데, 회사 개수가 얼마 되지 않아 플로이드 워셜 알고리즘을 사용하면 좀 더 간단하게 풀 수 있을 것 같다.