LeetCode Solutions
787. Cheapest Flights Within K Stops
Time: $O(|E| + |V|\log |E|)$ Space: $O(nk)$
class Solution {
public:
int findCheapestPrice(int n, vector<vector<int>>& flights, int src, int dst,
int k) {
vector<vector<pair<int, int>>> graph(n);
using T = tuple<int, int, int>; // (d, u, stops)
priority_queue<T, vector<T>, greater<>> minHeap;
vector<vector<int>> dist(n, vector<int>(k + 2, INT_MAX));
minHeap.emplace(0, src, k + 1); // Start with node src with d == 0
dist[src][k + 1] = 0;
for (const vector<int>& f : flights) {
const int u = f[0];
const int v = f[1];
const int w = f[2];
graph[u].emplace_back(v, w);
}
while (!minHeap.empty()) {
const auto [d, u, stops] = minHeap.top();
minHeap.pop();
if (u == dst)
return d;
if (stops > 0)
for (const auto& [v, w] : graph[u]) {
const int newDist = d + w;
if (newDist < dist[v][stops - 1]) {
dist[v][stops - 1] = newDist;
minHeap.emplace(newDist, v, stops - 1);
}
}
}
return -1;
}
};
class Solution {
public int findCheapestPrice(int n, int[][] flights, int src, int dst, int k) {
List<Pair<Integer, Integer>>[] graph = new List[n];
// (d, u, stops)
Queue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
int[][] dist = new int[n][k + 2];
Arrays.stream(dist).forEach(A -> Arrays.fill(A, Integer.MAX_VALUE));
for (int i = 0; i < n; ++i)
graph[i] = new ArrayList<>();
for (int[] f : flights) {
final int u = f[0];
final int v = f[1];
final int w = f[2];
graph[u].add(new Pair<>(v, w));
}
minHeap.offer(new int[] {0, src, k + 1}); // Start with node src with d == 0
dist[src][k + 1] = 0;
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)
for (Pair<Integer, Integer> node : graph[u]) {
final int v = node.getKey();
final int w = node.getValue();
final int newDist = d + w;
if (newDist < dist[v][stops - 1]) {
dist[v][stops - 1] = newDist;
minHeap.offer(new int[] {d + w, v, stops - 1});
}
}
}
return -1;
}
}
class Solution:
def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
graph = [[] for _ in range(n)]
minHeap = [(0, src, k + 1)] # (d, u, stops)
dist = [[math.inf] * (k + 2) for _ in range(n)]
for u, v, w in flights:
graph[u].append((v, w))
while minHeap:
d, u, stops = heapq.heappop(minHeap)
if u == dst:
return d
if stops > 0:
for v, w in graph[u]:
newDist = d + w
if newDist < dist[v][stops - 1]:
dist[v][stops - 1] = newDist
heapq.heappush(minHeap, (newDist, v, stops - 1))
return -1