rustgym/leetcode/
_666_path_sum_4.rs

1struct Solution;
2use std::collections::HashMap;
3
4impl Solution {
5    fn path_sum(nums: Vec<i32>) -> i32 {
6        let mut tree: HashMap<(usize, usize), i32> = HashMap::new();
7        for x in nums {
8            let v = x % 10;
9            let p = ((x / 10) % 10 - 1) as usize;
10            let d = ((x / 100) - 1) as usize;
11            *tree.entry((d, p)).or_default() = v;
12        }
13        let mut res = 0;
14        let mut path: Vec<i32> = vec![];
15        Self::dfs((0, 0), &mut path, &mut res, &tree);
16        res
17    }
18
19    fn dfs(
20        start: (usize, usize),
21        path: &mut Vec<i32>,
22        sum: &mut i32,
23        tree: &HashMap<(usize, usize), i32>,
24    ) {
25        let val = tree[&start];
26        path.push(val);
27        let left = (start.0 + 1, start.1 * 2);
28        let right = (start.0 + 1, start.1 * 2 + 1);
29        if !tree.contains_key(&left) && !tree.contains_key(&right) {
30            *sum += path.iter().copied().sum::<i32>();
31        }
32        if tree.contains_key(&left) {
33            Self::dfs(left, path, sum, tree);
34        }
35        if tree.contains_key(&right) {
36            Self::dfs(right, path, sum, tree);
37        }
38        path.pop();
39    }
40}
41
42#[test]
43fn test() {
44    let nums = vec![113, 215, 221];
45    let res = 12;
46    assert_eq!(Solution::path_sum(nums), res);
47    let nums = vec![113, 221];
48    let res = 4;
49    assert_eq!(Solution::path_sum(nums), res);
50}