666. Path Sum IV

If the depth of a tree is smaller than 5, then this tree can be represented by a list of three-digits integers.

For each integer in this list:

  1. The hundreds digit represents the depth D of this node, 1 <= D <= 4.
  2. The tens digit represents the position P of this node in the level it belongs to, 1 <= P <= 8. The position is the same as that in a full binary tree.
  3. The units digit represents the value V of this node, 0 <= V <= 9.

Given a list of ascending three-digits integers representing a binary tree with the depth smaller than 5, you need to return the sum of all paths from the root towards the leaves.

It's guaranteed that the given list represents a valid connected binary tree.

Example 1:

Input: [113, 215, 221]
Output: 12
Explanation: 
The tree that the list represents is:
    3
   / \
  5   1

The path sum is (3 + 5) + (3 + 1) = 12.

 

Example 2:

Input: [113, 221]
Output: 4
Explanation: 
The tree that the list represents is: 
    3
     \
      1

The path sum is (3 + 1) = 4.

Rust Solution

struct Solution;
use std::collections::HashMap;

impl Solution {
    fn path_sum(nums: Vec<i32>) -> i32 {
        let mut tree: HashMap<(usize, usize), i32> = HashMap::new();
        for x in nums {
            let v = x % 10;
            let p = ((x / 10) % 10 - 1) as usize;
            let d = ((x / 100) - 1) as usize;
            *tree.entry((d, p)).or_default() = v;
        }
        let mut res = 0;
        let mut path: Vec<i32> = vec![];
        Self::dfs((0, 0), &mut path, &mut res, &tree);
        res
    }

    fn dfs(
        start: (usize, usize),
        path: &mut Vec<i32>,
        sum: &mut i32,
        tree: &HashMap<(usize, usize), i32>,
    ) {
        let val = tree[&start];
        path.push(val);
        let left = (start.0 + 1, start.1 * 2);
        let right = (start.0 + 1, start.1 * 2 + 1);
        if !tree.contains_key(&left) && !tree.contains_key(&right) {
            *sum += path.iter().copied().sum::<i32>();
        }
        if tree.contains_key(&left) {
            Self::dfs(left, path, sum, tree);
        }
        if tree.contains_key(&right) {
            Self::dfs(right, path, sum, tree);
        }
        path.pop();
    }
}

#[test]
fn test() {
    let nums = vec![113, 215, 221];
    let res = 12;
    assert_eq!(Solution::path_sum(nums), res);
    let nums = vec![113, 221];
    let res = 4;
    assert_eq!(Solution::path_sum(nums), res);
}

Having problems with this solution? Click here to submit an issue on github.