[알고리즘] 파이썬으로 플로이드-워셜 알고리즘 구현하기 | 모든 노드의 최단거리를 구할 수 있는 쉽지만 계산이 많은 알고리즘

플로이드-워셜(Floyd-Warshall) 알고리즘은 특정 노드가 아닌 모든 노드에서 최단 경로를 찾을 수 있는 알고리즘 입니다. 그래프의 모든 노드에서 다른 노드까지 최단 경로를 계산하는 것으로 (특정 시점에 국한된 것이 아닌) 각 노드에서 다른 노드까지 최단 경로의 행렬을 반환합니다.

 

3개의 for문으로 비교적 쉽게 구할 수 있지만, 그만큼 많은 연산을 요구합니다. 코딩 테스트 문제에서 만약 나온다고 하면, 비교적 노드 범위 개수가 적게 주어질 경우 이 알고리즘을 적용해보는 것도 방법이 될 듯 합니다. 

1 단계 : 초기화

D[s][e]는 노드 s에서 e까지 가는 최단 거리를 저장하는 거리 리스트 입니다. 먼저 이 리스트를 출발 ~ 도착 지점이 같은 경우 (s =e) 0, 다른 칸을 무한대의 값으로 초기화 합니다. 

 

V = len(graph) # 노드 개수 
D = [[float('inf') * (V+1) for _ in range(V+1)]]

# 자기 자신에게 가는 거리 = 0
for i in range(1, V+1): 
    D[i][i] = 0

2 단계 : 그래프 데이터 저장

출발 노드[s]와 도착 노드[e], 가중치 w라고 했을 때 D[s][e] = w 로 저장을 진행합니다. 코드에서는 graph를 따로 입력 받은 것을 가정하기 때문에 아래와 같이 코드를 입력할 수 있습니다. 

 

# 그래프 데이터 저장
for i in range(V):
    for j in range(V):
        D[i][j] = graph[i][j]

 

3 단계 : 점화식 

플로이드-워셜 알고리즘은 아래 점화식을 진행하면 됩니다. 

 

 

위 점화식은 아래와 같이 구현 가능합니다. 

 

# 점화식에 따라 업데이트
for k in range(1, V+1):
    for s in range(1, V+1):
        for e in range(1, V+1):
            if D[s][e] > (D[s][k] + D[k][e]):
                D[s][e] = (D[s][k] + D[k][e])

 

최종 코드

위 단계를 모두 종합하면 아래와 같이 요약될 수 있습니다. 아래 코드에서 반환하는 D 행렬은 출발 ~ 도착 노드에서 최단 거리를 구할 수 있습니다. 예를 들어, 1부터 5까지 가는 최단 거리는 D[1][5]를 통해 구할 수 있습니다.

 

def floyd_warshall(graph):
    V = len(graph) # 노드 개수 
    D = [[float('inf') * (V+1) for _ in range(V+1)]]
    
    for i in range(1, V+1): 
        D[i][i] = 0

    for i in range(V):
        for j in range(V):
            D[i][j] = graph[i][j]
    
    for k in range(1, V+1):
        for s in range(1, V+1):
            for e in range(1, V+1):
                if D[s][e] > (D[s][k] + D[k][e]):
                    D[s][e] = (D[s][k] + D[k][e])

    return D