2714. Find Shortest Path with K Hops

• Time: $O((|V| + |E|)\log |V|)$
• Space: $O(|E| + |V|)$

C++JavaPython

class Solution {
 public:
  // Similar to 787. Cheapest Flights Within K Stops
  int shortestPathWithHops(int n, vector<vector<int>>& edges, int s, int d,
                           int k) {
    vector<vector<pair<int, int>>> graph(n);

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

    return dijkstra(graph, s, d, 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 + 1, INT_MAX));
    using T = tuple<int, int, int>;  // (d, u, hops)
    priority_queue<T, vector<T>, greater<>> minHeap;

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

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

    throw;
  }
};

class Solution {
  // Similar to 787. Cheapest Flights Within K Stops
  public int shortestPathWithHops(int n, int[][] edges, int s, int d, int k) {
    List<Pair<Integer, Integer>>[] graph = new List[n];

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

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

    return dijkstra(graph, s, d, k);
  }

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

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

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

    throw new IllegalArgumentException();
  }
}

class Solution:
  # Similar to 787. Cheapest Flights Within K Stops
  def shortestPathWithHops(self, n: int, edges: List[List[int]], s: int, d: int, k: int) -> int:
    graph = [[] for _ in range(n)]

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

    return self._dijkstra(graph, s, d, k)

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

    dist[src][k] = 0
    minHeap = [(dist[src][k], src, k)]  # (d, u, hops)

    while minHeap:
      d, u, hops = heapq.heappop(minHeap)
      if u == dst:
        return d
      if dist[u][hops] > d:
        continue
      for v, w in graph[u]:
        # Go from u -> v with w cost.
        if d + w < dist[v][hops]:
          dist[v][hops] = d + w
          heapq.heappush(minHeap, (dist[v][hops], v, hops))
        # Hop from u -> v with 0 cost.
        if hops > 0 and d < dist[v][hops - 1]:
          dist[v][hops - 1] = d
          heapq.heappush(minHeap, (dist[v][hops - 1], v, hops - 1))