Skip to content

787. Cheapest Flights Within K Stops 👍

  • Time: $O((|V| + |E|)\log |V|)$
  • Space: $O(|E| + |V|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
 public:
  int findCheapestPrice(int n, vector<vector<int>>& flights, int src, int dst,
                        int k) {
    vector<vector<pair<int, int>>> graph(n);

    for (const vector<int>& flight : flights) {
      const int u = flight[0];
      const int v = flight[1];
      const int w = flight[2];
      graph[u].emplace_back(v, w);
    }

    return dijkstra(graph, src, dst, k);
  }

 private:
  int dijkstra(const vector<vector<pair<int, int>>>& graph, int src, int dst,
               int k) {
    vector<vector<int>> dist(graph.size(), vector<int>(k + 2, INT_MAX));
    using T = tuple<int, int, int>;  // (d, u, stops)
    priority_queue<T, vector<T>, greater<>> minHeap;

    dist[src][k + 1] = 0;
    minHeap.emplace(dist[src][k + 1], src, k + 1);

    while (!minHeap.empty()) {
      const auto [d, u, stops] = minHeap.top();
      minHeap.pop();
      if (u == dst)
        return d;
      if (stops == 0 || d > dist[u][stops])
        continue;
      for (const auto& [v, w] : graph[u])
        if (d + w < dist[v][stops - 1]) {
          dist[v][stops - 1] = d + w;
          minHeap.emplace(dist[v][stops - 1], v, stops - 1);
        }
    }

    return -1;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
  public int findCheapestPrice(int n, int[][] flights, int src, int dst, int k) {
    List<Pair<Integer, Integer>>[] graph = new List[n];

    for (int i = 0; i < n; i++)
      graph[i] = new ArrayList<>();

    for (int[] flight : flights) {
      final int u = flight[0];
      final int v = flight[1];
      final int w = flight[2];
      graph[u].add(new Pair<>(v, w));
    }

    return dijkstra(graph, src, dst, k);
  }

  private int dijkstra(List<Pair<Integer, Integer>>[] graph, int src, int dst, int k) {
    int[][] dist = new int[graph.length][k + 2];
    Arrays.stream(dist).forEach(A -> Arrays.fill(A, Integer.MAX_VALUE));
    // (d, u, stops)
    Queue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]);

    dist[src][k + 1] = 0;
    minHeap.offer(new int[] {dist[src][k + 1], src, k + 1});

    while (!minHeap.isEmpty()) {
      final int d = minHeap.peek()[0];
      final int u = minHeap.peek()[1];
      final int stops = minHeap.poll()[2];
      if (u == dst)
        return d;
      if (stops == 0 || d > dist[u][stops])
        continue;
      for (Pair<Integer, Integer> pair : graph[u]) {
        final int v = pair.getKey();
        final int w = pair.getValue();
        if (d + w < dist[v][stops - 1]) {
          dist[v][stops - 1] = d + w;
          minHeap.offer(new int[] {dist[v][stops - 1], v, stops - 1});
        }
      }
    }

    return -1;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution:
  def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
    graph = [[] for _ in range(n)]

    for u, v, w in flights:
      graph[u].append((v, w))

    return self._dijkstra(graph, src, dst, k)

  def _dijkstra(self, graph: List[List[Tuple[int, int]]], src: int, dst: int, k: int) -> int:
    dist = [[math.inf for _ in range(k + 2)] for _ in range(len(graph))]

    dist[src][k + 1] = 0
    minHeap = [(dist[src][k + 1], src, k + 1)]  # (d, u, stops)

    while minHeap:
      d, u, stops = heapq.heappop(minHeap)
      if u == dst:
        return d
      if stops == 0 or d > dist[u][stops]:
        continue
      for v, w in graph[u]:
        if d + w < dist[v][stops - 1]:
          dist[v][stops - 1] = d + w
          heapq.heappush(minHeap, (dist[v][stops - 1], v, stops - 1))

    return -1