[알고리즘] 플로이드 워셜 예제

이것이 취업을 위한 코딩 테스트다. with 파이썬 책을 참고하여 정리한 내용입니다.

📖 미래 도시

공중 미래 도시에는 1번부터 N번 까지의 회사가 있는데, 특정 회사끼리는 서로 도로를 통해 연결되어 있다. 도로는 양방향으로 이동할 수 있으며, 이동하는 데에는 정확히 1만큼의 시간으로 이동할 수 있다. 이때 1번 회사에서 K번 회사를 방문한 다음 X번 회사로 가는 최소 시간을 계산하는 프로그램을 작성하시오.

✔️ 입력 조건

• 첫째 줄에 전체 회사의 개수 N과 경로의 개수 M이 공백으로 구분되어 차례대로 주어진다. (1 $\leq$ N, M $\leq$ 100)이 주어진다.
• 둘째 줄부터 M + 1번째 줄에는 연결된 두 회사의 번호가 공백으로 구분되어 주어진다.
• M + 2번째 줄에는 X와 K가 공백으로 구분되어 차례대로 주어진다. (1 $\leq$ K $\leq$ 100)

✔️ 출력 조건

• 첫째 줄에 방문 판매원 A가 K번 회사를 거쳐 X번 회사로 가는 최소 이동 시간을 출력한다.
• 만약 X번 회사에 도달할 수 없다면 -1을 출력한다.

✔️ 입력 예시 1

5 7
1 2
1 3
1 4
2 4
3 4
3 5
4 5
4 5

✔️ 출력 예시 1

3

✔️ 입력 예시 2

4 2
1 3
2 4
3 4

✔️ 출력 예시 2

-1

🗒️ 문제 해설

이 문제는 전형적인 플로이드 워셜 알고리즘 문제이다. 현재 문제에서 N의 범위가 100이하로 매우 한정적이다. 따라서 구현이 간단한 플로이드 워셜 알고리즘을 사용하는 것이 유리하다.

import sys

input = sys.stdin.readline
INF = int(1e9)

n, m = map(int, input().split())
# 2차원 리스트(그래프 표현)를 만들고, 모든 값을 무한으로 초기화
graph = [[INF] * (n + 1) for _ in range(n + 1)]

# 자기 자신으로의 비용은 0
for i in range(n + 1):
    for j in range(n + 1):
        if i == j:
            graph[i][j] = 0
for _ in range(m):
    # a <-> b 의 비용 1 (양방향)
    a, b = map(int, input().split())
    graph[a][b] = 1
    graph[b][a] = 1

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

for i 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][i] + graph[i][b])

result = graph[1][k] + graph[k][x]
if result >= INF:
    print(-1)
else:
    print(result)