P131 二叉树中和为给定值的路径

描述#

这个题目说的是，给你一棵二叉树和一个整数，你要找到这棵二叉树上从根到叶子节点路径和等于这个整数的所有路径。

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

/   

2     4
\   / 

10 8  16
给你的整数是 13。
在这棵二叉树中存在两条从根到叶子节点的路径，它们的路径和等于 13。分别是：
[
[1, 2, 10],
[1, 4, 8]
]
因此你要返回以上两条路径。

代码#

public class AlgoCasts {

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

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

  private void path(TreeNode root, int sum,
                    List<Integer> elem, List<List<Integer>> result) {
    if (root == null) return;
    elem.add(root.val);
    if (root.left == null && root.right == null && root.val == sum)
      result.add(new ArrayList<>(elem));
    path(root.left, sum - root.val, elem, result);
    path(root.right, sum - root.val, elem, result);
    elem.remove(elem.size()-1); // T: O(1)
  }

  // Time: O(n), Space: O(n)
  public List<List<Integer>> pathSumRecursive(TreeNode root, int sum) {
    List<List<Integer>> result = new ArrayList<>();
    List<Integer> elem = new ArrayList<>();
    path(root, sum, elem, result);
    return result;
  }

  // Time: O(n), Space: O(n)
  public List<List<Integer>> pathSumIterative(TreeNode root, int sum) {
    List<List<Integer>> result = new ArrayList<>();
    List<Integer> elem = new ArrayList<>();
    Stack<TreeNode> stack = new Stack<>();
    Set<TreeNode> visited = new HashSet<>();
    int curSum = 0;

    while (root != null || !stack.isEmpty()) {
      while (root != null) {
        elem.add(root.val);
        curSum += root.val;
        visited.add(root);
        stack.push(root);
        root = root.left;
      }
      TreeNode n = stack.peek();
      if (n.right == null || visited.contains(n.right)) {
        if (n.left == null && n.right == null && curSum == sum)
          result.add(new ArrayList<>(elem));
        stack.pop();
        elem.remove(elem.size()-1);
        curSum -= n.val;
        root = null;
      } else root = n.right;
    }
    return result;
  }

  // Time: O(n), Space: O(n)
  public List<List<Integer>> pathSumIterativePrev(TreeNode root, int sum) {
    List<List<Integer>> result = new ArrayList<>();
    List<Integer> elem = new ArrayList<>();
    Stack<TreeNode> stack = new Stack<>();
    TreeNode prev = null;
    int curSum = 0;

    while (root != null || !stack.isEmpty()) {
      while (root != null) {
        elem.add(root.val);
        curSum += root.val;
        stack.push(root);
        root = root.left;
      }
      TreeNode n = stack.peek();
      if (n.right == null || n.right == prev) {
        if (n.left == null && n.right == null && curSum == sum)
          result.add(new ArrayList<>(elem));
        stack.pop();
        elem.remove(elem.size()-1);
        curSum -= n.val;
        prev = n;
        root = null;
      } else root = n.right;
    }
    return result;
  }

}