Medium | LeetCode 437. 路径总和 III | 树 + 回溯 + 前缀和

437. 路径总和 III

给定一个二叉树，它的每个结点都存放着一个整数值。

找出路径和等于给定数值的路径总数。

路径不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。

二叉树不超过1000个节点，且节点数值范围是 [-1000000,1000000] 的整数。

示例：

root = [10,5,-3,3,2,null,11,3,-2,null,1], sum = 8

      10
     /  \
    5   -3
   / \    \
  3   2   11
 / \   \
3  -2   1

返回 3。和等于 8 的路径有:

1.  5 -> 3
2.  5 -> 2 -> 1
3.  -3 -> 11

解题思路

如果把这道题改为只能从根节点出发, 那么这道题只需要使用先序遍历的方法遍历一遍树即可。如下

int res = 0;

public void pathSum(TreeNode root, int target) {
    if (target == 0) {
        res++;
    }
    if (root == null) {
        return;
    }
    pathSum(root.left, target - root.val);
    pathSum(root.right, target - root.val);
}

那这道题它可以从任意节点出发, 我们想到只要把访问的路径保存起来, 然后访问一个节点时, 栈里的数字是否能够组合成目标值。但是这样时间复杂度会比较高。

其实可以借鉴 Easy | LeetCode 1. 两数之和 | 排序+双指针 | HaspMap 的思路, 判断当前的节点是否能作为一个满足条件的路径的最后一个节点, 只需要判断, 是否有一条路径, 从根节点出发, 然后在此路径上, 存在target - root.val 的路径和。这个时候, 很自然想到, 在先序遍历的同时, 需要将从根节点到当前节点的路径的所有节点的前缀和。

private int res = 0;

public int pathSum(TreeNode root, int sum) {
    // 记录从根节点 到 当前节点的路径的总和
    int curSum = 0;
    // 用来保存从根节点到当前节点的路径上所有节点的前缀和
    Map<Integer, Integer> prefixSum = new HashMap<>();
    // 初始化一个路径长度为0的前缀和路径
    prefixSum.put(0, 1);
    // 先序遍历树
    preOrder(root, sum, curSum, prefixSum);
    return res;
}

private void preOrder(TreeNode root, int sum, int curSum, Map<Integer, Integer> prefixSum) {
    if (root == null) {
        return;
    }
    // 计算从根节点到当前节点的和
    curSum += root.val;
    // 然后判断是否是否存在前缀和 为curSum - sum(目标值) 的前缀和存在
    res += prefixSum.getOrDefault(curSum - sum, 0);
    // 把当前节点计算进前缀和里, 准备向左右子树找
    prefixSum.put(curSum, prefixSum.getOrDefault(curSum, 0) + 1);
    // 看左子树是否存在符合条件的路径和
    preOrder(root.left, sum, curSum, prefixSum);
    // 看右子树是否存在符合条件的路径和左
    preOrder(root.right, sum, curSum, prefixSum);
    // 左右子树都判断完了, 然后回溯, 需要将当前节点的路径和移除
    int curSumCount;
    if ((curSumCount = (prefixSum.get(curSum) - 1)) > 0) {
        // 如果移除之后此路径和个数不是0, 大于0, 则还是保留这个路径和
        prefixSum.put(curSum, curSumCount);
    } else {
        // 否则, 移除这个路径和的KEY
        prefixSum.remove(curSum);
    }

}

方法二：暴力遍历

以任意一个节点作为起始节点, 判断是否存在从起始节点, 到其子孙节点的路径和为SUM的路径。

public int pathSum(TreeNode root, int sum) {
    return root == null ? 
        0 : 
    // 以当前节点为根节点的的路径和 + 左孩子为根节点的路径和 + 右孩子为根节点的路径和
    (preOrder(root, sum, 0) + pathSum(root.left, sum) + pathSum(root.right, sum));
}
private int preOrder(TreeNode root, int sum, int curSum) {
    return root == null ? 
        0 : 
    	(((curSum = (curSum + root.val)) == sum ? 1 : 0) + 
     		preOrder(root.left, sum, curSum) + 
     		preOrder(root.right, sum, curSum));
}