[이것이 코딩 테스트다] with Python

[이것이 코딩 테스트다] 다익스트라 알고리즘 with Python

Evolving Developer 2023. 5. 28. 13:52

🔒다익스트라 최단 경로 알고리즘

💡 그래프에서 여러 개의 노드가 있을 때,

💡 특정한 노드에서 출발하여 다른 노드로 가는 각각의 최단 경로를 구해주는 알고리즘

🔒특이점

  • 음의 간선이 없을 때 정상 작동
    • 음의 간선이란 0보다 작은 값을 가지는 간선
    • 현실엔 음의 간선이 존재하지 않기 때문에 실제로 GPS 소프트웨어의 기본 알고리즘으로 채택됨
  • 다익스트라 알고리즘은 기본적으로 그리디 알고리즘으로 분류됨
  • 각 노드에 대한 현재까지의 최단 거리 정보를 항상 1차원 배열에 저장하며, 리스트를 계속 갱신함

알고리즘

💡 1. 출발 노드를 설정

💡 2. 최단 거리 테이블을 초기화

💡 3. 방문하지 않은 노드 중에서 최단 거리가 가장 짧은 노드를 선택

💡 4. 해당 노드를 거쳐 다른 노드로 가는 비용을 계산하여 최단 거리 테이블 갱신

💡 5. 위 과정에서 3번, 4번을 반복

다익스트라 알고리즘을 구현하는 두 가지 방법

💡 방법 1. 구현하기 쉽지만 느리게 동작하는 코드

💡방법 2. 구현하기에 조금 더 까다롭지만 빠르게 동작하는 코드

방법 1

  • 이미 방문해 탐색을 마친 노드를 다시 방문하지 않기 위해 방문 여부를 저장하는 visited[] 배열 선언
    • 처음에는 그 어떤 노드도 방문하지 않았단 의미에서 False로 초기화
    • visited[0]을 제외한 visited[1] ~ visited[n]까지의 정보를 저장해야 하기 때문에 (n+1) 크기의 배열 생성
    • visited = [False] * (n+1)
  • ✅방문하지 않은 노드 중에서, 최단 거리가 가장 짧은 노드를 찾기 위해 get_smallest_node() 함수 선언
    • 최단 거리를 저장하는 distance[] 배열에서 최단 거리를 찾고, visited[] 배열 값이 False인지 확인 후 반환
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
  • ✅get_smallest_node() 함수로 최단 거리가 가장 짧은 노드를 찾았다면, 이제 그 노드와 연결된 노드와 그 거리를 탐색하기 시작
    • get_smallest_node()가 반환한 노드가 3번 노드이고, 3번 노드와 연결된 노드가 2, 5번 노드라면
    • 시작 노드에서 2번, 5번 노드로 바로 가는 게 빠른지
    • 1번 노드에서 3번 노드를 경유해 2, 5번 노드로 가는 게 빠른지 비교 후
    • 더 작은 값으로 distance[2]와 distance[5] 업데이트
    • 위 과정을 반복하며 더 짧은 경로를 찾을 때마다 distance[] 업데이트

💻 코드

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

# 노드의 개수, 간선의 개수 입력받기
n, m = map(int, input().split())
# 노드의 시작 번호 입력받기
start = int(input())

# 노드 정보를 담을 그래프 만들기
graph = [[] for _ in range(n+1)]
# 방문한 적이 있는지 체크하는 리스트
visited = [False] * (n+1)
# 최단 거리를 모두 무한으로 초기화
distance = [INF] * (n+1)

# 노드와 간선 정보 입력받기
for _ in range(m):
    node1, node2, dis = map(int, input().split())
    graph[node1].append((node2, dis))

# 방문하지 않은 노드 중에서, 가장 최단 거리가 짧은 노드의 번호를 반환
# 반환된 노드를 기준으로 연결된 노드를 재탐색해야 하기 때문
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):
    # 시작 노드와 시작 노드 사이 거리는 0
    distance[start] = 0
    # 시작 노드와 연결된 노드와의 거리로 distance[] 초기화
    for node in graph[start]:
        distance[node[0]] = node[1]

    # 시작 노드를 제외한 (n-1)개의 노드에 대해서
    # 탐색하며 최단 경로를 업데이트 해야 하기 때문에 아래 과정을 (n-1)번 반복
    for i in range(n):
        # 방문하지 않았으며, 시작 노드에서 경로가 가장 짧은 노드 선택
        now = get_smallest_node()
        visited[now] = True
        # 선택 노드 기준으로 연결된 노드의 최단 거리를 탐색
        for node in graph[now]:
            # node[0] 노드까지 선택 노드를 경유해서 가는 최단 거리
            cost = distance[now] + node[1]
            # 경유해서 가는 방법이 최단 거리라면
            if cost < distance[node[0]]:
                # 거리 저장 배열을 수정
                distance[node[0]] = cost

dijkstra(start)

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

방법 2

  • ✅최단 거리가 가장 짧은 노드를 반환하는 get_smallest_node() 함수 대신 우선순위 큐 사용
    • 우선순위 큐란 우선순위를 가진 데이터들을 저장하는 큐를 의미
    • 데이터를 꺼낼 때 우선 순위가 높은 데이터가 가장 먼저 나온다는 특징
    • heapq.heappush(배열, (거리, 노드))
    • 위와 같이 저장하고 heapq.heappop()을 하면 자동적으로 거리가 가장 짧은 노드와 거리를 반환해준다.
  • ✅이미 방문한 노드는 제외하기 위해 visited[] 배열 대신 아래 조건문 사용
# 방문할 노드를 가져오되, 이미 방문한 노드는 제외
dist, now = heapq.heappop(q)
# 방문할 노드
if distance[now] < dist:
    continue
  • 어떻게 위와 같은 수식이 나올까? 예시를 들어보자.
  • 만약 다음과 같은 그래프가 있다면
  • 1 --1--> 2
  • 1 --9--> 3
  • 2 --1--> 3
  • 시작 노드는 1이며, 노드 1과 노드 1의 사이 거리는 0이다.
    • 우선순위 큐: (0, 1)
    • heappop() -> (0, 1)
  • 노드 1과 연결된 노드 2와 노드 3의 정보가 거리가 짧은 순으로 우선순위 큐에 담긴다.
    • 우선순위 큐: (1, 2) (9, 3)
    • heappop() -> (1, 2)
    • distance[2] = 1, distance[3] = 9
  • 노드 2와 연결된 노드 3의 정보가 우선순위 큐에 담기고, 노드 3으로 가는 최단 거리가 2로 수정된다.
    • 우선순위 큐: (2, 3) (9, 3)
    • heappop() -> (2, 3)
    • distance[3] = 2
  • 노드 3과 연결된 노드는 없으므로 그 어떤 것도 큐에 담기지 않는다.
    • 우선순위 큐: (9, 3)
    • heappop() -> (9, 3)
  • 이때 distance[3] = 2이고, dist = 9이므로 distance[3] < dist를 충족하므로, 방문된 노드라고 판단하고 넘어간다.

💻 코드

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

# 노드의 개수, 간선의 개수 입력받기
n, m = map(int, input().split())
# 노드의 시작 번호 입력받기
start = int(input())

# 노드 정보를 담을 그래프 만들기
graph = [[] for _ in range(n+1)]
# 최단 거리를 모두 무한으로 초기화
distance = [INF] * (n+1)

# 노드와 간선 정보 입력받기
for _ in range(m):
    node1, node2, dis = map(int, input().split())
    graph[node1].append((node2, dis))

def dijkstra(start):
    # 방문할 노드를 우선순위 큐로 관리
    q = []
    # 시작 노드와 시작 노드 사이 거리는 0
    distance[start] = 0
    # 우선순위 큐에 시작 노드를 담고 시작
    heapq.heappush(q, (0, start))

    # 더이상 방문할 노드가 없을 때까지 반복
    while q:
        # 방문할 노드를 가져오되, 이미 방문한 노드는 제외
        dist, now = heapq.heappop(q)
        # 방문할 노드
        if distance[now] < dist:
            continue
        # 방문할 노드와 연관된 노드 정보 가져오기
        for node in graph[now]:
            # node[0]은 인덱스, node[1]은 거리
            # cost = 현재 노드까지의 최단 거리 + 새 노드로의 이동 거리
            cost = dist + node[1]
            # 만약 현재 노드를 경유해서 가는 방법이 기존 방법보다 더 최단 거리라면
            if cost < distance[node[0]]:
                # 최단 거리 업데이트
                distance[node[0]] = cost
                # 우선순위 큐에 연관 노드와 (예상) 최단 거리 삽입
                heapq.heappush(q, (cost, node[0]))

dijkstra(start)

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