Question

666. 路径总和 IV

English Version

题目描述

对于一棵深度小于 5 的树，可以用一组三位十进制整数来表示。对于每个整数：

• 百位上的数字表示这个节点的深度 d，1 <= d <= 4。
• 十位上的数字表示这个节点在当前层所在的位置 P， 1 <= p <= 8。位置编号与一棵满二叉树的位置编号相同。
• 个位上的数字表示这个节点的权值 v，0 <= v <= 9。

给定一个包含三位整数的 升序 数组 nums ，表示一棵深度小于 5 的二叉树，请你返回 _从根到所有叶子结点的路径之和 _。

保证 给定的数组表示一个有效的连接二叉树。

 

示例 1：

输入: nums = [113, 215, 221]
输出: 12
解释: 列表所表示的树如上所示。
路径和 = (3 + 5) + (3 + 1) = 12.

示例 2：

输入: nums = [113, 221]
输出: 4
解释: 列表所表示的树如上所示。
路径和 = (3 + 1) = 4.

 

提示:

• 1 <= nums.length <= 15
• 110 <= nums[i] <= 489
• nums 表示深度小于 5 的有效二叉树

解法

方法一

class Solution:
  def pathSum(self, nums: List[int]) -> int:
    def dfs(node, t):
      if node not in mp:
        return
      t += mp[node]
      d, p = divmod(node, 10)
      l = (d + 1) * 10 + (p * 2) - 1
      r = l + 1
      nonlocal ans
      if l not in mp and r not in mp:
        ans += t
        return
      dfs(l, t)
      dfs(r, t)

    ans = 0
    mp = {num // 10: num % 10 for num in nums}
    dfs(11, 0)
    return ans

class Solution {
  private int ans;
  private Map<Integer, Integer> mp;

  public int pathSum(int[] nums) {
    ans = 0;
    mp = new HashMap<>(nums.length);
    for (int num : nums) {
      mp.put(num / 10, num % 10);
    }
    dfs(11, 0);
    return ans;
  }

  private void dfs(int node, int t) {
    if (!mp.containsKey(node)) {
      return;
    }
    t += mp.get(node);
    int d = node / 10, p = node % 10;
    int l = (d + 1) * 10 + (p * 2) - 1;
    int r = l + 1;
    if (!mp.containsKey(l) && !mp.containsKey(r)) {
      ans += t;
      return;
    }
    dfs(l, t);
    dfs(r, t);
  }
}

class Solution {
public:
  int ans;
  unordered_map<int, int> mp;

  int pathSum(vector<int>& nums) {
    ans = 0;
    mp.clear();
    for (int num : nums) mp[num / 10] = num % 10;
    dfs(11, 0);
    return ans;
  }

  void dfs(int node, int t) {
    if (!mp.count(node)) return;
    t += mp[node];
    int d = node / 10, p = node % 10;
    int l = (d + 1) * 10 + (p * 2) - 1;
    int r = l + 1;
    if (!mp.count(l) && !mp.count(r)) {
      ans += t;
      return;
    }
    dfs(l, t);
    dfs(r, t);
  }
};

func pathSum(nums []int) int {
  ans := 0
  mp := make(map[int]int)
  for _, num := range nums {
    mp[num/10] = num % 10
  }
  var dfs func(node, t int)
  dfs = func(node, t int) {
    if v, ok := mp[node]; ok {
      t += v
      d, p := node/10, node%10
      l := (d+1)*10 + (p * 2) - 1
      r := l + 1
      if _, ok1 := mp[l]; !ok1 {
        if _, ok2 := mp[r]; !ok2 {
          ans += t
          return
        }
      }
      dfs(l, t)
      dfs(r, t)
    }
  }
  dfs(11, 0)
  return ans
}