leetcode

0 class Solution {

1 const int mod = 1000000007;

2

3 struct edge {

4 int dest, w;

5 edge(int d, int w) : dest(d), w(w) {}

6 friend bool operator<(const edge &e1, const edge &e2) { return e1.w > e2.w; }

7 };

8

9 vector<int> dijkstra(int n, vector<vector<edge>> &adj, int start) {

10 vector<int> dist(n + 1, INT_MAX);

11 priority_queue<edge> pq;

12

13 vector<int> visited(n + 1, false);

14 pq.push({n, dist[n] = 0});

15 for (int k = 0; k < n; k++) {

16 while (!pq.empty() && visited[pq.top().dest])

17 pq.pop();

18 if (pq.empty()) break;

19 edge c = pq.top();

20 pq.pop();

21 visited[c.dest] = true;

22 for (edge &p : adj[c.dest])

23 if (!visited[p.dest] && dist[c.dest] + p.w < dist[p.dest]) {

24 dist[p.dest] = dist[c.dest] + p.w;

25 pq.push({p.dest, dist[p.dest]});

26 }

27 }

28 return dist;

29 }

30

31 public:

32 int countRestrictedPaths(int n, vector<vector<int>> &edges) {

33 vector<vector<edge>> adj(n + 1);

34

35 for (int i = 0; i < edges.size(); i++) {

36 adj[edges[i][0]].push_back({edges[i][1], edges[i][2]});

37 adj[edges[i][1]].push_back({edges[i][0], edges[i][2]});

38 }

39

40 vector<int> dist = dijkstra(n, adj, n);

41 vector<int> dp(n + 1, 0);

42 vector<int> visited(n + 1, 0);

43 stack<int> st;

44

45 int count = 0;

46 st.push(1);

47 while (!st.empty()) {

48 int root = st.top();

49

50 if (root == n) {

51 st.pop();

52 count++;

53 continue;

54 }

55

56 if (root == -1) {

57 st.pop();

58 root = st.top();

59 st.pop();

60 dp[root] += count;

61 continue;

62 }

63

64 dp[root] = -count;

65 visited[root] = true;

66 st.push(-1);

67 for (auto [c, w] : adj[root])

68 if (dist[root] > dist[c]) {

69 if (visited[c] && dp[c] >= 0)

70 count = (count + dp[c]) % mod;

71 else

72 st.push(c);

73 }

74 }

75

76 return count;

77 }

78 };