P143 二叉树中和为给定值的路径数量

描述#

这个题目说的是，给你一棵二叉树和一个目标值，你要计算出路径和等于目标值的路径一共有多少条。其中，路径不需要开始于根节点，或结束于叶子节点。但必须是从父节点到子节点。

比如说，给你的二叉树是：
1

/   

2     4
\   / 

3 5   6
给你的目标值是 5。
在这棵二叉树中，路径和等于 5 的路径有：
1 -> 4
2 -> 3
5
总共有 3 条满足条件的路径，因此你要返回 3。

代码#

public class AlgoCasts {

  public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int x) {
      val = x;
    }
  }

  private int pathFrom(TreeNode root, int sum) {
    if (root == null) return 0;
    int cnt = 0;
    if (root.val == sum) ++cnt;
    cnt += pathFrom(root.left, sum - root.val);
    cnt += pathFrom(root.right, sum - root.val);
    return cnt;
  }

  // Time: O(n^2), Space: O(n)
  public int pathSumRecursive(TreeNode root, int sum) {
    if (root == null) return 0;
    return pathFrom(root, sum) +
      pathSumRecursive(root.left, sum) +
      pathSumRecursive(root.right, sum);
  }

  private int dfs(TreeNode root, int cur, int sum,
                  Map<Integer, Integer> prefixSum) {
    if (root == null) return 0;
    cur += root.val;
    int cnt = prefixSum.getOrDefault(cur - sum, 0);
    prefixSum.put(cur, prefixSum.getOrDefault(cur, 0) + 1);
    cnt += dfs(root.left, cur, sum, prefixSum);
    cnt += dfs(root.right, cur, sum, prefixSum);
    prefixSum.put(cur, prefixSum.get(cur) - 1);
    return cnt;
  }

  // Time: O(n), Space: O(n)
  public int pathSumPrefixSum(TreeNode root, int sum) {
    Map<Integer, Integer> prefixSum = new HashMap<>();
    prefixSum.put(0, 1);
    return dfs(root, 0, sum, prefixSum);
  }

}