リンク

問題の説明

入力

しゅつりょく

Input: n = 4, flights = [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]], src = 0, dst = 3, k = 1
Output: 700
Explanation:
The graph is shown above.
The optimal path with at most 1 stop from city 0 to 3 is marked in red and has cost 100 + 600 = 700.
Note that the path through cities [0,1,2,3] is cheaper but is invalid because it uses 2 stops.

せいげんじょうけん

1 <= n <= 100
0 <= flights.length <= (n * (n - 1) / 2)
flights[i].length == 3
0 <= fromi, toi < n
fromi != toi
1 <= pricei <= 104
There will not be any multiple flights between two cities.
0 <= src, dst, k < n
src != dst

完全なコード

import heapq

class Solution:
    
    def dijkstra(self,graph:dict,start:int,end:int,k:int)->int:
        heap=[]
        visited_set=set()
        # (cost,node,visit_count)
        visited_set.add((0,start,0))
        
        #(node,cost,visit_count)
        heapq.heappush(heap,(0,start,0))
        
        while heap:
            cost,node,visit_count=heapq.heappop(heap)
            if node==end:
                return cost
        
            for adj_node,adj_cost in graph[node]:
                new_state=(cost+adj_cost,adj_node,visit_count+1)
                if visit_count+1<=k and new_state not in visited_set:
                    visited_set.add(new_state)
                    heapq.heappush(heap,new_state)
                
        return -1
            
        
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph=dict()
        for i in range(n):
            graph[i]=[]
            
        for u,v,cost in flights:
            graph[u].append((v,cost))
        
        return self.dijkstra(graph,src,dst,k+1)
        

解説