Reference: LeetCode
Difficulty: Medium

Problem

Given a binary tree and a sum, find all root-to-leaf paths where each path’s sum equals the given sum.

Example: 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
Given:
the below binary tree and sum = 22,
      5
     / \
    4   8
   /   / \
  11  13  4
 /  \    / \
7    2  5   1
Output:
[
   [5,4,11,2],
   [5,8,4,5]
]

Analysis

Methods:

  1. Recursion (Preorder traversal)
  • Time: $O(N + Lh)$ where $L$ is the number of valid paths. 1 ms
    • Visiting all nodes take $O(N)$.
    • Creating $L$ valid paths, each one takes $h$ time to create.
  • Space: $O(h + Lh) = O(Lh)$
  • Note: We can make a new copy of list in each function call, but it takes a lot of space.
  1. Iteration
  • I tried preorder, inorder, bfs. All of them are difficult to implement because of the timing of recovering the path list.
  • Finally, I add a map depthMap to record the depth of each node.
    Time: $O(N + Lh)$. And recovering the path list takes time. 6 ms
    Space: $O(N)$

Code

Recursion

Note:

  • Remember to recover the path.
  • Learn the way to make a new copy of list.
  • Be careful of the generic type:
    • Correct: LinkedList<List<Integer>> result = new LinkedList<>();
    • Incorrect: LinkedList<LinkedList<Integer>> result = new LinkedList<>();
      • The return type is List<List<Integer>>, which is not compatible with LinkedList<LinkedList<Integer>>.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
public List<List<Integer>> pathSum(TreeNode root, int sum) {
  LinkedList<Integer> path = new LinkedList<>();
  LinkedList<List<Integer>> result = new LinkedList<>();
  pathSum(root, sum, path, result);
  return result;
}

private void pathSum(TreeNode root, int sum, LinkedList<Integer> path, LinkedList<List<Integer>> result) {
  if (root == null) {
    return;
  }
  
  path.addLast(root.val); // update the path O(1)
  
  if (root.left == null && root.right == null) { // a leaf
    if (root.val == sum) {
      result.add(new LinkedList(path)); // O(h)
      // could not return here, because we need to restore the path
    }
  } else {
    pathSum(root.left, sum - root.val, path, result);
    pathSum(root.right, sum - root.val, path, result);
  }
  
  path.removeLast();  // recover the path O(1)
}

Iteration

Note:

  • The order! Push into stack the right child first.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
public List<List<Integer>> pathSum(TreeNode root, int sum) {
  LinkedList<Integer> path = new LinkedList<>();
  LinkedList<List<Integer>> result = new LinkedList<>();
  HashMap<TreeNode, Integer> depthMap = new HashMap<>();
  if (root == null) {
    return result;
  }
  // init
  Stack<TreeNode> stack = new Stack<>();
  Stack<Integer> sums = new Stack<>();
  stack.push(root); depthMap.put(root, 1);
  sums.push(sum);

  while (stack.size() > 0) {
    TreeNode p = stack.pop();
    int depth = depthMap.get(p);
    int s = sums.pop();
    path.addLast(p.val);
    
    if (p.left == null && p.right == null) { // leaf
      if (p.val == s) result.add(new LinkedList<>(path));
    }
    
    // order!
    if (p.right != null) { // right
      stack.push(p.right); depthMap.put(p.right, depth + 1);
      sums.push(s - p.val);
    }
    
    if (p.left != null) { // left
      stack.push(p.left); depthMap.put(p.left, depth + 1);
      sums.push(s - p.val);
    }
    
    if (stack.size() > 0) {
      int d = depthMap.get(stack.peek());
      while (path.size() > d - 1) {
        path.removeLast();
      }
    }
  }

  return result;
}