Number of Restricted Paths From First to Last Node LeetCode solution in Rust

1
use std::collections::{BinaryHeap, HashMap, HashSet};
2

3
impl Solution {
4
  pub fn find_restricted(
5
    node: usize,
6
    shortest_path: &Vec<i32>,
7
    graph: &Vec<Vec<(usize, i32)>>,
8
    memo: &mut HashMap<usize, i32>,
9
  ) -> i32 {
10
    let last_node = shortest_path.len() - 1;
11
    if node == last_node {
12
      return 1;
13
    }
14

15
    if memo.contains_key(&node) {
16
      return *memo.get(&node).unwrap();
17
    }
18

19
    let mut pathes = 0;
20
    let MOD = i32::pow(10, 9) + 7;
21

22
    for edge in &graph[node] {
23
      let neighbor = edge.0;
24
      if shortest_path[neighbor] >= shortest_path[node] && neighbor != last_node {
25
        continue;
26
      }
27
      pathes += Solution::find_restricted(neighbor, shortest_path, graph, memo);
28
      pathes %= MOD;
29
    }
30

31
    memo.insert(node, pathes);
32
    return pathes;
33
  }
34

35
  // O(n) time,
36
  // O(n) space,
37
  // Approach: djikstra shortest path, dp, top down dp, heap,
38
  pub fn count_restricted_paths(n: i32, edges: Vec<Vec<i32>>) -> i32 {
39
    let mut graph = vec![Vec::new(); n as usize + 1];
40

41
    for edge in edges {
42
      let node1 = edge[0] as usize;
43
      let node2 = edge[1] as usize;
44
      let weight = edge[2];
45

46
      graph[node1].push((node2, weight));
47
      graph[node2].push((node1, weight));
48
    }
49

50
    let mut shortest_dist = vec![0; n as usize + 1];
51
    let mut min_heap = BinaryHeap::new();
52
    let mut visited: HashSet<usize> = HashSet::new();
53

54
    min_heap.push((0, n as usize));
55
    // calulate shortest path using djikstra
56
    while min_heap.len() > 0 {
57
      let popped = min_heap.pop().unwrap();
58
      let dist = -popped.0;
59
      let node = popped.1 as usize;
60

61
      if visited.contains(&node) {
62
        continue;
63
      }
64
      visited.insert(node);
65
      shortest_dist[node] = dist;
66

67
      for edge in &graph[node] {
68
        if visited.contains(&edge.0) {
69
          continue;
70
        }
71
        let new_dist = -(dist + edge.1);
72
        min_heap.push((new_dist, edge.0));
73
      }
74
    }
75

76
    return Solution::find_restricted(1, &shortest_dist, &graph, &mut HashMap::new());
77
  }
78
}