[알고리즘] 다익스트라 알고리즘 (+LeetCode 1514)

다익스트라 (Dijkstra) 알고리즘은 weighted graph에서 최단 경로를 찾는 알고리즘이다. weight가 양수일 때만 사용 가능하다.

 

여러 버전이 있는데 우선 위키피디아에 가장 먼저 뜨는 알고리즘을 살펴보자. weighted graph와 출발 노드(node)가 주어졌다고 하자. 처음에는 현제 노드가 출발 노드와 같고, 출발 노드에서 다른 모든 노드까지 거리를 무한대로 초기화한다. 단 출발 노드에서 출발 노드까지 거리는 0으로 설정한다. 다익스트라 알고리즘은 다음과 같다:

1. 방문하지 않은 노드 중 시작점과 거리가 가장 짧은 노드를 현재 노드로 설정
2. 인근 노드까지 출발 점에서부터 거리를 계산한다
3. 이렇게 계산한 거리가 이전에 구했던 출발점부터의 거리보다 더 작다면 새로 계산한 값으로 업데이트한다

참고로 이 알고리즘은 한번 방문한 노드까지의 거리가 최단으로 설정된다. 하지만 노드 사이에 음수 weight가 존재하면 더 이상 최단이지 않을 수 있기 때문에 작동하지 않는다. 

 

pseudo code는 다음과 같다 (pseudo code에서는 경로도 구하기 때문에 이전 노드 정보를 prev에 저장한다):

 1  function Dijkstra(Graph, source):
 2      
 3      for each vertex v in Graph.Vertices:
 4          dist[v] ← INFINITY # v에서부터 거리를 무한으로 일단 설정
 5          prev[v] ← UNDEFINED # v까지 최단 경로 (v의 이전 위치). 우선 미정
 6          add v to Q # 모든 vertex를 일단 queue Q에 넣음.
 7      dist[source] ← 0 # source까지 거리는 0
 8      
 9      while Q is not empty:
10          u ← vertex in Q with min dist[u] # u 까지 최단 거리인 vertex를 가져옴. 처음에는 source.
11          remove u from Q # queue에서 u 제거
12          
13          for each neighbor v of u still in Q: # u의 neighbor 중 아직 방문하지 않은 vertex v가 있으면 (Q에 있다는 뜻)
14              alt ← dist[u] + Graph.Edges(u, v) # vertex까지 거리 계산
15              if alt < dist[v]: # v까지 거리가 예전에 구한 거리보다 작으면
16                  dist[v] ← alt # 최단 거리 값 업데이트
17                  prev[v] ← u # 최단 거리 경로 업데이트
18
19      return dist[], prev[]

만약 특정 조건을 충족하는 node를 찾는 것이라면 10번째 줄에서 u를 받아서 탈출 조건을 확인하면 된다.

 

이것을 그대로 구현해도 괜찮지만 더 깔끔하게도 구현 가능하다 (대신 알고리즘이 살짝 바뀐다). 밑에 파이썬 코드에서 graph는 {node:(neighbor_node,neighbor_dist)} 꼴이라고 가정했다. 참고로 dist와 prev를 합쳐서 distNnode라는 dictionary에 저장했다. 프린트하는 코드가 없으면 10줄정도밖에 안 된다.

 1    def dijkstra(graph, start):
 2      distNnode = collections.defaultdict(tuple) # {그래프의 노드:(거리와,이전 노드)}
 3      q = [(0,start,None)] # (거리, 노드, 이전 노드)
 4      idx = 0 # 프린트 용도
 5      while q:
 6          dist, node,prev_node = heapq.heappop(q)
 7          if node not in distNnode: # 한번만 방문하게 된다
 8              print('-----',idx,'-----')
 9              print('q:',q)
10              print(f'dist:{dist},node:{node},prev_node:{prev_node}')
11              idx+=1
12              distNnode[node] = (dist,prev_node)
13              for neighbor_node, neighbor_dist in graph[node]:
14                  heapq.heappush(q,(dist + neighbor_dist,neighbor_node,node))
15      return distNnode

다음과 같이 작동한다.

1. min heap에 우선 시작 점을 추가한다. (코드 line 3) 
2. min heap이 빌 때까지 반복:
   1. min heap에서 시작 점부터 최단 거리인 노드를 받아온다 (코드 line 6)
   2. 이 노드를 방문한 적 없으면 (코드 line 7) 인근 노드까지 거리를 계산해서 min heap에 추가한다

이 알고리즘에서는 방문하는 노드에서 인근 노드까지 거리를 heap에 모두 추가한다. 하지만 방문했던 노드를 필터링하기 때문에 결국 모든 노드를 한 번만 방문하게 된다 (연결된 그래프인 경우). 또한 처음 소개한 알고리즘 line 15에서 인근 노드까지 최단 거리를 업데이트하는 코드가 있는데, min heap을 사용하기 때문에 자동으로 이런 부분들이 처리된다.

 

다음 그래프에 대해서 위 알고리즘을 한번 돌려보자:

graph = { 0:[(1,1),(2,4),], 1:[(0,1),(4,2),], 2:[(0,4),(3,1),(4,3),(5,2),], 3:[(2,1),], 4:[(1,2),(2,3),(5,5),], 5:[(4,5),(2,2),] }

출력 값

0 -> 1 -> 4 -> 2 -> 3 -> 5 순으로 방문하는 것을 알 수 있다. 출력 값 마지막 줄은 {노드 : (최단 거리, 이전 노드)} 꼴이다.

 

이제 문제를 하나 풀어보자

Path with Maximum Probability - LeetCode

 

 

 

 

출발 점과 끝 점이 주어졌을 때 갈 수 있는 경로 중 가장 높은 확률로 도달할 수 있는 경로를 찾는 문제이다. 경로의 확률은 edge에 주어진 weight의 곱들로 계산한다. 

다익스트라 알고리즘을 거의 그대로 사용하면 된다. 단, 파이썬의 heapq는 min heap이고 구하고 싶은 것은 max 확률이기 때문에, 값을 (1-확률) 값을 사용하면 된다.

class Solution:
    def maxProbability(self, n: int, edges: List[List[int]], succProb: List[float], start: int, end: int) -> float:
        # 우선 graph부터 만들자
        graph = collections.defaultdict(list)
        for edge,prob in zip(edges,succProb):
            graph[edge[0]].append((edge[1],prob))
            graph[edge[1]].append((edge[0],prob))
        
        # 경로는 (이전 node) 저장할 필요 없기 때문에 확률 값만 저장
        prob_dict = collections.defaultdict(float) 
        
        q = [(0,start)] # min heap -> store 1-prob (i.e. reversed probability) in queue; prob in prob_dict
        while q:
            rev_prob, node = heapq.heappop(q) 
            if node not in prob_dict:
                prob_dict[node] = 1-rev_prob # rev_prob을 확률로 변환
                for neighbor_node, prob in graph[node]:
                    neighbor_prob = (1-rev_prob) * prob # 경로 확률 계산
                    heapq.heappush(q,(1-neighbor_prob,neighbor_node))
        return prob_dict[end]

Refs:

https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm

박성길, 정진호 "파이썬 알고리즘 인터뷰: 95가지 알고리즘 문제 풀이로 완성하는 코딩 테스트"