Number of Ways to Arrive at Destination LeetCode solution in Rust

1
use std::collections::*;
2

3
const MOD: usize = 1_000_000_007;
4
const INF: usize = 1_000_000_000_000_000;
5
fn dfs(memo: &Vec<(usize, HashSet<usize>)>, dp: &mut Vec<usize>, i: usize) -> usize {
6
  if dp[i] == INF {
7
    let mut temp = 0;
8
    for &j in &memo[i].1 {
9
      temp += dfs(memo, dp, j);
10
      temp %= MOD;
11
    }
12
    dp[i] = temp;
13
  }
14

15
  dp[i]
16
}
17

18
impl Solution {
19
  pub fn count_paths(n: i32, roads: Vec<Vec<i32>>) -> i32 {
20
    let n = n as usize;
21
    let mut g = vec![vec![]; n];
22
    for arr in roads {
23
      let a = arr[0] as usize;
24
      let b = arr[1] as usize;
25
      let c = arr[2] as usize;
26
      g[a].push((b, c));
27
      g[b].push((a, c));
28
    }
29

30
    let mut memo = vec![(INF, HashSet::new()); n];
31
    memo[0].0 = 0;
32
    let mut stack = vec![(0, 0)];
33
    while !stack.is_empty() {
34
      let mut new_stack = vec![];
35
      while let Some((ci, cv)) = stack.pop() {
36
        for &(ni, v) in &g[ci] {
37
          let nv = cv + v;
38
          if nv < memo[ni].0 {
39
            memo[ni].1 = HashSet::new();
40
            memo[ni].1.insert(ci);
41
            memo[ni].0 = nv;
42
            new_stack.push((ni, nv));
43
          } else if nv == memo[ni].0 {
44
            memo[ni].1.insert(ci);
45
          }
46
        }
47
      }
48
      stack = new_stack;
49
    }
50

51
    let mut result = vec![INF; n];
52
    result[0] = 1;
53

54
    dfs(&memo, &mut result, n - 1) as i32
55
  }
56
}