본문 바로가기
알고리즘 배우기

[알고리즘 정리] 다익스트라

by cwin 2024. 7. 14.
import heapq

def dijkstra(adj_matrix, start):
    num_nodes = len(adj_matrix)
    # 각 노드까지의 최단 거리 초기화
    distances = [float('infinity')] * num_nodes
    distances[start] = 0

    # 우선순위 큐 초기화
    priority_queue = [(0, start)]  # (거리, 노드)

    while priority_queue:
        current_distance, current_node = heapq.heappop(priority_queue)

        # 현재 노드의 거리가 저장된 거리보다 크면 무시
        if current_distance > distances[current_node]:
            continue

        # 인접 노드 확인
        for neighbor in range(num_nodes):
            weight = adj_matrix[current_node][neighbor]
            if weight > 0:  # 0보다 큰 경우에만 유효한 간선
                distance = current_distance + weight

                # 더 짧은 경로가 발견되면 업데이트
                if distance < distances[neighbor]:
                    distances[neighbor] = distance
                    heapq.heappush(priority_queue, (distance, neighbor))

    return distances

# 예시 인접 행렬 (그래프)
# 0: 노드가 연결되지 않음
adj_matrix = [
    [0, 1, 4, 0, 0],
    [1, 0, 2, 5, 0],
    [4, 2, 0, 1, 0],
    [0, 5, 1, 0, 3],
    [0, 0, 0, 3, 0]
]

# 다익스트라 알고리즘 실행
start_node = 0  # 시작 노드 (인덱스)
shortest_paths = dijkstra(adj_matrix, start_node)
print(shortest_paths)  # 결과: [0, 1, 3, 2, 6]