[알고리즘 algorithm] 다익스트라 알고리즘- 파이썬 코드와 예제

최단 경로 문제를 푸는 알고리즘으로 가장 잘 알려진 다익스트라 알고리즘을 파이썬 코드로 작성하면 다음과 같다.

 

 

기본 문제 : leetcode.com/problems/network-delay-time/

Network Delay Time - LeetCode

class Solution:
    def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
        graph = collections.defaultdict(list)
        
        # u,v,w = departure, arrive, time
        for u,v,w in times:
            graph[u].append([v,w])
            
        
        dist = collections.defaultdict(int)
            
        # time, departure
        Q= [[0,k]]
        
        while Q:
            time, node = heapq.heappop(Q)
            if node not in dist:
                dist[node] = time
                for v, w in graph[node]:
                    alt = time + w
                    heapq.heappush(Q, [alt, v])
                    
        if len(dist) == n:
            return max(dist.values())
        
        return -1

다익스트라 알고리즘은 BFS 탐색을 기본으로 하며, 최소 힙을 사용해 구현한다.

 

 

 

응용 문제 : leetcode.com/problems/cheapest-flights-within-k-stops

Cheapest Flights Within K Stops - LeetCode

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        
        graph = collections.defaultdict(list)
        
        for u,v,w in flights:
            graph[u].append([v,w])
          
        # stops, cost
        result = []
        
        # cost, node, stops
        Q = [[0,src,-1]]
        
        while Q:
            cost, node, stops = heapq.heappop(Q)
            if node == dst :
                return cost
            if stops < K:
                for v,w in graph[node]:
                    alt = cost + w
                    heapq.heappush(Q, [alt,v,stops+1])
                    
        return -1

출처:

<파이썬 알고리즘 인터뷰> 박상길 지음, 정진호 일러스트, 책만, 2020(https://www.onlybook.co.kr/entry/algorithm-interview)