파이썬 다익스트라 (Dijkstra) 알고리즘

본문은  '이것이 코딩테스트다' 책(나동빈 저, 한빛미디어)을 참고하며 작성하였습니다.

파이썬 다익스트라 (Dijkstra) 알고리즘 (혹은 데이크스트라 알고리즘)

 

: 그래프에서 한 정점에서 다른 여러 개의 정점들 간의 최단 경로를 찾는 알고리즘

(그리디 알고리즘으로도 분류된다. - 가장 비용이 적게 드는 정점을 선택해서 과정을 반복하기 때문)

 

이러한 그래프가 있을 때 1번에서 다른 정점으로 가는 최단 경로를 모두 구해보는 문제를 생각했을 때,

 

1. 초기화 - 모든 거리를 무한대로 설정한다. (파이썬에서는 통상적으로 1e9 사용) --> 2 (무한) , 3 (무한), 4 (무한), 5 (무한), 6 (무한) 
2. 1번에서 갈 수 있는 정점들 확인 후 비용 계산 후 최소값으로 갱신 --> 2 (2) , 3 (5), 4 (1), 5 (무한), 6 (무한)
3.  방문하지 않은 정점 중에서 최단 거리가 가장 짧은 정점 선택 (여기서는 4번 선택), 4번에서 갈 수 있는 정점 확인 후 비용 계산. 4번에서는 3번, 5번으로 갈 수 있다. 1->4->3 일 때 비용은 4로, 앞선 단계에서 구한 비용보다 적게 드므로 갱신, 1->4->5 일 때는 비용이 2이므로 갱신 -->   2 (2) , 3 (4), 4 (1), 5 (2), 6 (무한)
4. 앞선 단계를 반복한다. 다음으로 최단 거리가 짧은 정점은 2와 5인데 보통 번호가 적은 정점을 선택한다.
5. 최종적으로 2 (2) , 3 (4), 4 (1), 5 (2), 6 (4) 로 각각의 최단 경로를 계산할 수 있다.

이 과정을 요약해서 말하면 '방문하지 않은 정점 중에서 가장 최단 거리가 짧은 정점을 선택' 하는 과정을 반복하는 것이다.

계산 과정 3번 이후 한 번 선택된 정점은 다시 갱신되지 않는다. 즉, 이 알고리즘은 한 단계 당 하나의 정점에 대한 최단 거리를 확실히 찾는 것!

 

이제 파이썬으로 위의 알고리즘을 작성해보자.

 

[방법 1] - 시간 복잡도 O(V^2) (V: # of nodes)

* 인덱싱 하기 쉽게 n 대신 n+1로 graph를 만든다.

1. 초기화

import sys

# input = sys.stdin.readline
inf = int(1e9)

# example
n, m = 6, 11 # number of node and line
start = 1 # start number

# A list of information about nodes connected to each node.
graph = [ [] for i in range(n+1)]  


visited = [False] * (n+1)
distance = [inf] * (n+1)

2. 모든 간선 정보 입력받기

for _ in range(m):
	a, b, c = map(int, input().split()) # c is the cost for a -> b

graph[a].append((b,c))

3. 방문하지 않은 노드 중에서, 가장 최단 거리가 짧은 노드의 번호를 반환하는 함수

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

4. 다익스트라 함수 만들기

def dijkstra(start):
    # initialization
    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

방법 1은 간단한 다익스트라 알고리즘을 구현하는 방법이지만, 노드의 개수가 많아지면 문제를 해결하기 어렵다.

최단 거리가 가장 짧은 노드를 찾을 때 일일이 모든 원소를 하나씩 탐색했지만 (선형 시간 소요), 힙(Heap) 자료 구조를 이용하면 로그 시간이 소요된다!  (힙에 대한 포스팅은 따로 할 예정.. )

 

[방법 2] - 시간 복잡도 O(ElogV) (V: # of nodes, E: # of lines)

방법 2를 위의 예시에 다시 적용해보면, 

1. 초기화 - 모든 거리를 무한대로 설정한다. (파이썬에서는 통상적으로 1e9 사용) --> 2 (무한) , 3 (무한), 4 (무한), 5 (무한), 6 (무한) 
2. 우선순위 큐에 (거리, 노드) 순으로 튜플을 넣는다. 처음에는 (0, 1) 넣기. (1에서 시작하므로)
3.  우선순위 큐에서 원소를 꺼내고, 방문하지 않은 노드 처리 - 1번 노드와 이어져 있는 2, 3, 4번 노드부터 --> (1, 4) , (2, 2), (5, 3) ** 힙은 튜플의 맨 앞 원소 기준으로 우선순위 결정 **
4. 큐에서 첫 번째 원소 (1, 4)를 꺼낸다 - 4번 노드에 연결돼있는 3번과 5번의 경로를 구해서 큐에 넣어준다 (갱신). 각각 1+3, 1+1 --> (2, 2), (2, 5), (4, 3), (5, 3)
5. 큐에서 첫 번째 원소 (2, 2)를 꺼낸다 - 2번 노드에 연결돼있는 노드들 (3번,5번)의 거리가 갱신되지 않는다. --> (2, 5), (4, 3), (5, 3)
6. 큐에서 첫 번째 원소 (2, 5)를 꺼낸다 - 5번 노드에 연결돼있는 3번과 6번의 경로를 구해서 큐에 넣어준다. 각각 2+1, 2+2 --> (3, 3), (4, 3), (4, 6), (5, 3)
7. 큐에서 첫 번째 원소 (3, 3)을 꺼낸다. --> (4, 3), (4, 6), (5, 3)
8. 큐에서 첫 번째 원소 (4, 3)을 꺼낸다. 하지만 3번 노드는 이미 처리된 노드이므로 무시 --> (4, 6), (5, 3)
9. 큐에서 첫 번째 원소 (4, 6)을 꺼낸다. 마지막 남은 (5, 3)은 마찬가지로 무시.
10. 따라서 꺼낸 원소들을 보면 (거리, 노드) 로 볼 때 (0, 1), (1, 4), (2, 2), (2, 5), (3, 3), (4, 6) 으로 구할 수 있다!

후.. 너무 구구절절 썼지만 이렇게 안 쓰면 내가 나중에 이해 못할까봐 ... 다 썼다.

이제 이걸 코드로 구현해보자! 방법1에서 다익스트라 함수만 바꿔주면 된다.

import heapq

def dajkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
    	dist, now = heapq.heappop(q)
        if distance[now] < dist:
        	contiune
        for i in graph[now]:
        	cost = dist + i[1]
            if cost < distance[i[0]]:
            	distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

일단 개념적으로 이해는 갔는데.. 실제로 문제를 풀 수 있을진 잘 모르겠다 ㅜ.ㅜ

다음에 공부할 내용은 플로이드 워셜 알고리즘 ,,,, 알고리즘 공부는 너무 어려우어어어엉