Algorithm/개념 정리

[최단경로] 다익스트라

nayeoniee 2022. 8. 13. 00:45

다익스트라

특정 노드 → 모든 노드로 가는 최단 경로를 계산한다. 정확히 얘기하면 최단 경로는 아니고, (일반적으로 코딩 테스트에서는) 최단 경로의 길이를 계산한다.

매번 비용이 최소가 되는 경로를 선택하기 때문에 그리디 알고리즘으로 분류되기도 한다.

 

1. 간단한 다익스트라 알고리즘

위와 같이 가중치가 있는 방향 그래프에서 다익스트라 알고리즘을 사용해 최단 경로를 구해보자. 

1번 노드에서 시작해 모든 노드로 가는 최단 경로를 구하는 예시이다. 다음에 선택되는 노드를 노란색으로 표시했다.

  1. 각 노드까지의 거리는 무한대로 초기화하고, 1번 노드는 자기 자신이므로 0으로 초기화한다.
  2. 가장 가까운 1번 노드를 방문해 현재 거리보다 더 짧은 경로를 발견하면 업데이트한다. 다음에는 가장 가까운 노드를 방문하고, 거리가 같다면 일반적으로 노드 번호가 작은 노드를 방문한다.
  3. 모든 노드를 방문할 때까지 반복한다.
import sys
read = sys.stdin.readline
INF = float('inf')

# 입력
n, m = map(int, read().split())  # 노드와 간선의 개수 
start = int(read())  # 시작 노드 
graph = [[] for _ in range(n+1)]  # 노드 연결 정보를 담을 리스트
visited = [False] * (n+1)  # 노드 방문 정보 저장
distance = [INF] * (n+1)  # 거리는 무한대로 초기화

for _ in range(m):
    n1, n2, cost = map(int, read().split())  # n1 -> n2로 가는 비용이 cost
    graph[n1].append((n2, cost))

# 방문하지 않는 노드 중 가장 거리가 짧은 노드 반환 (순차 탐색)
def get_smallest_node():
    idx = 0
    min_val = INF
    for i in range(1, n+1):
        if distance[i] < min_val and not visited[i]:
            min_val = distance[i]
            idx = i
    return idx

def dijkstra(start):
    # 시작 노드 초기화
    distance[start] = 0
    visited[True]
    
    # 시작 노드와 연결된 노드의 거리 업데이트
    for node, cost in graph[start]:
        distance[node] = cost

    for i in range(n-1):
        # 현재 가장 가까운 노드를 선택
        now = get_smallest_node()
        visited[now] = True

        # 선택한 노드와 연결된 요소들 확인
        for node, cost in graph[now]:
            new_cost = distance[now] + cost
            if distance[node] > new_cost:
                distance[node] = new_cost

# main
dijkstra(start)
print(distance[1:])

 

이 방법은 매번 “최단 거리가 가장 짧은 노드”를 찾기 위해서 거리가 담긴 distance 리스트를 선형적으로 탐색한다. 따라서 V가 노드의 개수일 때 시간 복잡도는 $O(V^2)$ 이다. 

 

 

2. 개선된 다익스트라 알고리즘

  • 우선순위 큐인 heapq 라이브러리를 사용해 최소값을 찾으면 매번 “최단 거리가 가장 짧은 노드”를 찾기 위해서 선형 탐색하는 시간을 줄일 수 있다.
  • q 리스트를 추가적으로 사용해 매번 가장 가까운 노드를 찾고, 방문 여부를 저장하는 visited는 필요하지 않다.
  • 간선의 개수를 E, 노드의 개수를 V라 표현하면 매번 가장 가까운 노드를 선택하는데 $O(logV)$, 해당 연산을 간선의 개수만큼 수행하므로 전체 시간 복잡도는 $O(E*logV)$ 이다.
import sys
import heapq
read = sys.stdin.readline
INF = float('inf')

# 입력
n, m = map(int, read().split())  # 노드와 간선의 개수 
start = int(read())  # 시작 노드 
graph = [[] for _ in range(n+1)]  # 노드 연결 정보를 담을 리스트
distance = [INF] * (n+1)  # 거리는 무한대로 초기화

for _ in range(m):
    n1, n2, cost = map(int, read().split())  # n1 -> n2로 가는 비용이 cost
    graph[n1].append((n2, cost))

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 node, cost in graph[now]:
            new_cost = dist + cost
            if new_cost < distance[node]:
                distance[node] = new_cost
                heapq.heappush(q, (new_cost, node))
    
# main
dijkstra(start)
print(distance[1:])