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;
}

Comment