다익스트라(python)

in Algorithm on Algorithm, Dijkstra, Greedy, Python

참고 자료

이것이 취업을 위한 코딩 테스트다 with 파이썬(도서)

다익스트라

  • 그래프에서 여러 개의 노드가 있을 때, 특정한 노드에서 출발하여 다른 노드로 가는 각각의 최단 경로를 구해주는 알고리즘.
  • ‘각 노드에 대한 현재까지의 최단 거리’ 정보를 1차원 리스트에 저장하며 계속 갱신한다는 특징이 있다.
  1. 출발 노드를 설정
  2. 최단 거리 테이블을 초기화
  3. 방문하지 않은 노드 중 출발 노드와의 최단 거리가 가장 짧은 노드 선택
  4. 해당 노드를 거쳐 다른 노드로 가는 비용을 계산하여 최단 거리 테이블을 갱신
  5. 3, 4번을 반복

간단한 다익스트라 구현

  • for문을 사용해서 가장 최단 거리의 노드를 탐색
  • 시간 복잡도 : O(N2), N은 노드의 개수
  • 화살표가 존재하는 간선
import sys
input = sys.stdin.readline
INF = int(1e9) # 무한을 의미하는 값으로 10억 설정

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

start = int(input()) # 시작 노드

graph = [[] for i in range(n + 1)] # 각 노드에 연결되어 있는 노드에 대한 정보를 담은 리스트

visited = [False] * (n + 1) # 방문했는지 체크하는 리스트

distance = [INF] * (n + 1) # 최단 거리 테이블

for _ in range(m):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))

# 방문하지 않은 노드들 중 가장 최단 거리의 노드를 반환
def get_smallest_node():
    min_value = INF
    index = 0
    for i in range(1, n + 1):
        if distance[i] < min_value and not visited[i]:
            min_value = distance[i]
            index = i
    return index

def dijkstra(start):
    # 출발 노드 정보 초기화
    distance[start] = 0
    visited[start] = True
    # 출발 노드의 간선으로 연결된 노드들 정보 갱신
    for j in graph[start]:
        distance[j[0]] = j[1]
    # 출발 노드를 제외한 나머지 노드들 탐색
    for i in range(n - 1):
        now = get_smallest_node()
        visited[now] = True
        for j in graph[now]:
            cost = distance[now] + j[1]
            if cost < distance[j[0]]:
                distance[j[0]] = cost

dijkstra(start)

for i in range(1, n + 1):
    if distance[i] == INF:
        print("INFINITY")
    else:
        print(distance[i])

개선된 다익스트라 구현

  • 최소 힙(min heap)을 사용하여 가장 최단 노드를 탐색
  • 시간 복잡도 : O(Mlog(N)), M은 간선의 수, N은 노드의 수
  • 화살표가 존재하는 간선
import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)

n, m = map(int, input().split())

start = int(input())

graph = [[] for i in range(n + 1)]

distance = [INF] * (n + 1)

for _ in range(m):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))

def dijkstra(start):
    q = []
    # heapq는 가장 앞에 값을 기준으로 정렬하기 때문에 (비용, 노드) 순으로 넣어야 한다.
    heapq.heappush(q,(0, start))
    distance[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        # 거리가 최단 거리 테이블에 있는 값보다 작은 경우 갱신할 필요가 없다.
        if distance[now] < dist:
            continue
        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

dijkstra(start)

for i in range(1, n + 1):
    if distance[i] ==INF:
        print("INFINITY")
    else:
        print(distance[i])