787. Cheapest Flights Within K Stops | 1786

 1class Solution:
 2    def findCheapestPrice(
 3        self, n: int, flights: List[List[int]], src: int, dst: int, k: int
 4    ) -> int:
 5        # dijkstra with k stop restrictions
 6
 7        # construct graph
 8        graph = {}  # source: [[distance0, destination0], [distance1, destination1]]
 9        for source, destination, distance in flights:
10            if graph.get(source, -1) == -1:
11                graph[source] = [[distance, destination]]
12            else:
13                graph[source].append([distance, destination])
14            if graph.get(destination, -1) == -1:
15                graph[destination] = []
16
17        # dijkstra init
18        dist = {k: float("inf") for k in graph.keys()}
19        parent = {k: None for k in graph.keys()}
20        pq = [(0, src, -1)]  # distance, location, stop count
21        visited = set()
22
23        while pq:
24            cur_dist, cur_loc, cur_stop_count = heapq.heappop(pq)
25
26            if cur_loc in visited or cur_stop_count > k:  # k-stop constraint
27                continue
28            visited.add(cur_loc)
29
30            for child_dist, child_loc in graph[cur_loc]:
31                possible_better_dist = cur_dist + child_dist
32                if possible_better_dist < dist[child_loc]:  # cheaper price
33                    dist[child_loc] = possible_better_dist
34                    parent[child_loc] = cur_loc
35                    heapq.heappush(
36                        pq, (possible_better_dist, child_loc, cur_stop_count + 1)
37                    )
38
39        return -1 if dist[dst] > 10**9 else dist[dst]

 1class Solution:
 2    def findCheapestPrice(
 3        self, n: int, flights: List[List[int]], src: int, dst: int, k: int
 4    ) -> int:
 5        # dijkstra with k stop restrictions
 6
 7        # construct graph
 8        # graph = {}  # source: [[distance0, destination0], [distance1, destination1]]
 9        # for source, destination, distance in flights:
10        #     if graph.get(source, -1) == -1:
11        #         graph[source] = [[distance, destination]]
12        #     else:
13        #         graph[source].append([distance, destination])
14        #     if graph.get(destination, -1) == -1:
15        #         graph[destination] = []
16
17        # the above does not work because there are cities that have no flights
18        graph = {i: [] for i in range(n)}
19        for source, destination, distance in flights:
20            graph[source].append([distance, destination])
21
22        # dijkstra variant init
23        # dist[node][stops_used] = best cost to reach node using exactly stops_used edges
24        dist = {key: [float("inf")] * (k + 2) for key in graph.keys()}
25        # parent = {k: None for k in graph.keys()}
26        pq = [(0, src, 0)]  # distance, location, flight count
27
28        while pq:
29            cur_dist, cur_loc, cur_flight_count = heapq.heappop(pq)
30
31            if cur_flight_count > k + 1:
32                continue
33
34            if cur_loc == dst:
35                return cur_dist  # min heap principles
36
37            if cur_dist > dist[cur_loc][cur_flight_count]:
38                continue
39
40            if cur_flight_count == k + 1:
41                continue
42
43            for child_dist, child_loc in graph[cur_loc]:
44                possible_better_dist = cur_dist + child_dist
45                new_flight_count = cur_flight_count + 1
46                if (
47                    possible_better_dist < dist[child_loc][new_flight_count]
48                ):  # cheaper price
49                    dist[child_loc][new_flight_count] = possible_better_dist
50                    heapq.heappush(
51                        pq, (possible_better_dist, child_loc, new_flight_count)
52                    )
53
54        return -1