-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
find-shortest-path-with-k-hops.cpp
44 lines (41 loc) · 1.58 KB
/
find-shortest-path-with-k-hops.cpp
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
// Time: O(n * k + (e * k) * log(n * k))
// Space: O(n * k + e)
// dijkstra's algorithm
class Solution {
public:
int shortestPathWithHops(int n, vector<vector<int>>& edges, int s, int d, int k) {
vector<vector<pair<int, int>>> adj(n);
for (const auto& e : edges) {
adj[e[0]].emplace_back(e[1], e[2]);
adj[e[1]].emplace_back(e[0], e[2]);
}
const auto& modified_dijkstra = [&]() {
static const int INF = numeric_limits<int>::max();
vector<vector<int>> best(size(adj), vector<int>(k + 1, INF));
best[s][0] = 0;
priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<tuple<int, int, int>>> min_heap;
min_heap.emplace(best[s][0], s, 0);
while (!empty(min_heap)) {
const auto [curr, u, cnt] = min_heap.top(); min_heap.pop();
if (curr > best[u][cnt]) {
continue;
}
if (u == d) {
return curr;
}
for (const auto& [v, w] : adj[u]) {
if (w < best[v][cnt] - curr) {
best[v][cnt] = curr + w;
min_heap.emplace(best[v][cnt], v, cnt);
}
if (cnt + 1 <= k && curr < best[v][cnt + 1]) {
best[v][cnt + 1] = curr;
min_heap.emplace(best[v][cnt + 1], v, cnt + 1);
}
}
}
return -1;
};
return modified_dijkstra();
}
};