1786.cpp (2262B)
1 class Solution { 2 const int mod = 1000000007; 3 4 struct edge { 5 int dest, w; 6 edge(int d, int w) : dest(d), w(w) {} 7 friend bool operator<(const edge &e1, const edge &e2) { return e1.w > e2.w; } 8 }; 9 10 vector<int> dijkstra(int n, vector<vector<edge>> &adj, int start) { 11 vector<int> dist(n + 1, INT_MAX); 12 priority_queue<edge> pq; 13 14 vector<int> visited(n + 1, false); 15 pq.push({n, dist[n] = 0}); 16 for (int k = 0; k < n; k++) { 17 while (!pq.empty() && visited[pq.top().dest]) 18 pq.pop(); 19 if (pq.empty()) break; 20 edge c = pq.top(); 21 pq.pop(); 22 visited[c.dest] = true; 23 for (edge &p : adj[c.dest]) 24 if (!visited[p.dest] && dist[c.dest] + p.w < dist[p.dest]) { 25 dist[p.dest] = dist[c.dest] + p.w; 26 pq.push({p.dest, dist[p.dest]}); 27 } 28 } 29 return dist; 30 } 31 32 public: 33 int countRestrictedPaths(int n, vector<vector<int>> &edges) { 34 vector<vector<edge>> adj(n + 1); 35 36 for (int i = 0; i < edges.size(); i++) { 37 adj[edges[i][0]].push_back({edges[i][1], edges[i][2]}); 38 adj[edges[i][1]].push_back({edges[i][0], edges[i][2]}); 39 } 40 41 vector<int> dist = dijkstra(n, adj, n); 42 vector<int> dp(n + 1, 0); 43 vector<int> visited(n + 1, 0); 44 stack<int> st; 45 46 int count = 0; 47 st.push(1); 48 while (!st.empty()) { 49 int root = st.top(); 50 51 if (root == n) { 52 st.pop(); 53 count++; 54 continue; 55 } 56 57 if (root == -1) { 58 st.pop(); 59 root = st.top(); 60 st.pop(); 61 dp[root] += count; 62 continue; 63 } 64 65 dp[root] = -count; 66 visited[root] = true; 67 st.push(-1); 68 for (auto [c, w] : adj[root]) 69 if (dist[root] > dist[c]) { 70 if (visited[c] && dp[c] >= 0) 71 count = (count + dp[c]) % mod; 72 else 73 st.push(c); 74 } 75 } 76 77 return count; 78 } 79 };