[Python] 가중치 그래프에서의 최단 경로

YUKI KIM·2022년 1월 17일

알고리즘 공부하면서

목록 보기

4/4

오늘 가중치 그래프에서 최단 경로를 구하는 문제를 푸는데 내가 너무 버벅여서 제대로 공부하기로 했다. 대표적인 최단 경로 알고리즘인 다익스트라, 벨만포드, 플로이드 와샬 알고리즘에 대한 개념과 파이썬으로 구현한 코드를 정리하고자 한다.

우선 본격적으로 들어가기 전, 설명의 편의를 위해 아래와 같이 가정을 한다.

V = 4    # 4개의 노드
E = 6    # 6개의 간선
K = 1    # 1번 노드에서 시작
graph = [[] for _ in range(V + 1)]  # graph[i] = (j, k): i->j 가중치는 k
dist = [float("inf") for _ in range(V + 1)] # 최단 경로 저장

다익스트라 (1→N, 음수 X)

다익스트라 알고리즘은 특정 정점에서 나머지 정점으로의 최단 경로를 구하는 알고리즘이다. 가중치가 모두 양수일 때만 사용 가능하다.

개념

다익스트라는 특정 정점에서 가중치가 가장 작은 정점을 선택해야 하므로 우선순위 큐를 사용한다. 파이썬에서는 우선순위 큐를 heapq 라이브러리를 임포트해서 사용 가능하다.

heapq 라이브러리에 대해 간단히 설명하자면 아래와 같다. 자세한 설명은 이 문서를 참고 바람.

import heapq

heap = []  # 힙을 저장할 리스트 생성
heapq.heappush(heap, N)  # 대상 리스트(heap)에 원소(N) 추가 
N = heapq.heappop(heap)  # 가장 작은 원소를 삭제 후 그 값을 리턴

다익스트라는 특정 정점에서 인접한 정점을 가중치가 작은 순으로 큐에 저장해가며 거리가 짧은 경로를 먼저 추출한다.

코드

def dijkstra(K, graph, dist):
    dist[K] = 0
    heap = []
    heapq.heappush(heap, (0, K))

    while heap:
        wei, now = heapq.heappop(heap)
        if dist[now] < wei:   # 더 짧은 경로 알고 있으면 무시
            continue

        for i in graph[now]:  # 인접 노드 검사
            next_wei = wei + i[1]
            if dist[i[0]] > next_wei:
                dist[i[0]] = next_wei
                heapq.heappush(heap, (next_wei, i[0]))

    return dist[1:]

힙에 (거리, 노드)형식으로 저장하고 있기 때문에 거리순으로 최소힙이 구성된다. 그런데 이렇게 하면 동일한 노드라도 힙에 저장될 수있다. (ex. [(10, 2), (7, 2)]) 하지만 if dist[now] < wei:에서 저장되어 있는 거리와 추출된 거리를 비교하기 때문에 이러한 경우에 대해서도 커버가 된다.

벨만포드 (1→N, 음수 O)

벨만포드 알고리즘은 특정 정점에서 나머지 정점으로의 최단 경로를 구하는 알고리즘이다. 가중치에 음수가 있어도 사용 가능하다. 또한 음수 사이클 존재의 여부를 알 수 있다.

개념

벨만포드 알고리즘은 다익스트라와 다르게 모든 경우를 다 탐색하는 방식으로, 모든 간선들로 모든 정점들을 이을 수 있을 만큼 (최단 거리를 만들기 위해) 계속 반복하는 방식이다.

V개의 정점은 (V-1)번 반복하면 모든 정점을 이을 수 있다. 매 반복 마다 모든 간선을 확인하고, 거리 정보가 더 짧으면 거리 정보를 갱신한다. 그런데 V번째 반복에도 거리 값이 갱신된다면 이는 음수 사이클이 존재한다는 의미라고 볼 수 있다.

코드

def bellmanford(V, K, graph, dist):
    dist[K] = 0
    minusCycle = False

    for i in range(V):   # 정점 수 만큼 반복
        for j in graph[i]:      # 모든 간선 확인
            cost = j[1]
            if dist[j[0]] > dist[i] + cost:
                dist[j[0]] = dist[i] + cost
                if i == V - 1:  # (V-1)번 이후에도 갱신되면 음수 사이클
                    minusCycle = True
     
    return dist[1:] if minusCycle is False else "음수 사이클"

방문하지 않는 노드 중, 최단 거리인 노드만을 방문하는 다익스트라와 달리 매 반복마다 갈 수 있는 모든 간선을 확인하고 있다.

플로이드 와샬 (N→N)

플로이드 와샬은 모든 정점에서 나머지 정점으로 가는 최단 경로를 구하는 알고리즘이다. 플로이드 와샬은 따로 시작 정점이 주어지지 않기 때문에 앞선 두 알고리즘과 차이가 있다.

개념

앞서 설명한 다익스트라와 벨만포드와 다르게 다이나믹 프로그래밍 기법으로 구현한다. 점화식은 dist[i][j] = min(dist[i][j], dist[i][n] + dist[n][j])이며, 정점 i에서 정점 n을 거쳐서 정점 j로 갈 때, n을 거쳐 가는 것이 더 최단경로일 경우를 업데이트 하는 의미라고 할 수 있겠다.

코드

def floydwarshall(V, graph):
    for k in range(V):
        for i in range(V):
            for j in range(V):
                if i == j:
                    graph[i][j] = 0 # 제자리는 거리 0
                graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j])

    return graph[K - 1]


if __name__ == '__main__':
    V, E = map(int, input().strip().split())
    graph = [[float("inf") for _ in range(V)] for _ in range(V)]
    
    for _ in range(E):
        u, v, e = map(int, input().strip().split())
        graph[u - 1][v - 1] = e
        

    print(solution(V, E, K, graph))