Reference: LeetCode
Difficulty: Medium

Problem

Given a binary tree, flatten it to a linked list in-place.

Example: 

1
2
3
4
5
    1
   / \
  2   5
 / \   \
3   4   6

The flattened tree should look like:

1
2
3
4
5
6
7
8
9
10
11
1
 \
  2
   \
    3
     \
      4
       \
        5
         \
          6

Analysis

Preorder

1
2
3
4
5
    1
   / \
  2   5
 / \   \
3   4   6

The idea is to use preorder traversal. For example, when root = 1, we want to set 1.right as 2 and set 4.right as 5.

Notice that to know where 4 is during the traversal, we can use a prev node to keep track of the previous node. The previous node of 1 is 4.

Note: I know the code could be shorter, but I think shorter code is barely readable!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
private TreeNode prevNode = null;

public void flatten(TreeNode root) {
  preorder(root);
}

private void preorder(TreeNode root) {
  if (root == null) {
    return;
  }
  prevNode = root; // visit root
  
  preorder(root.left); // go left
  TreeNode prevTemp = prevNode;
  preorder(root.right); // go right
  
  TreeNode rightTemp = root.right; // since prevTemp could be root
  root.right = root.left;
  root.left = null;
  
  prevTemp.right = rightTemp;
  prevTemp.left = null;
}

Time: $O(N)$
Space: $O(h)$

Recursive Function

Note: Just a bit modification on the previous solution.

This idea is a more recursive idea. Let’s say we have a flatten() function. For node 1, we can flatten its two children 2 and 5. At that time, we don’t need to know the detail but we are sure that we will have:

1
2
3
4
5
6
7
    1              1  // root
   / \            / \
  2   5          2   5
 / \   \          \   \
3   4   6          3   6
                    \
                     4

Two subproblems have been flattened.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
private TreeNode prevNode = null;

public void flatten(TreeNode root) {
  if (root == null) {
    return;
  }
  prevNode = root; // visit root
  
  flatten(root.left); // go left (preorder-like traversal)
  TreeNode prevTemp = prevNode;
  flatten(root.right); // go right
  
  TreeNode rightTemp = root.right; // since prevTemp could be root
  root.right = root.left;
  root.left = null;
  
  prevTemp.right = rightTemp;
  prevTemp.left = null;
}

Time: $O(N)$
Space: $O(N)$

LeetCode Solution

More like reversed postorder!

I don’t like this following method neither.

1
2
3
4
5
6
7
8
9
10
11
12
private TreeNode prev = null;

public void flatten(TreeNode root) {
  if (root == null) {
      return;
  }
  flatten(root.right);
  flatten(root.left);
  root.right = prev;
  root.left = null;
  prev = root;
}

Other Solution: (not using field variable)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
public void flatten(TreeNode root) {
  flatten(root, null);
}

private TreeNode flatten(TreeNode root, TreeNode prev) {
  if (root == null) {
    return prev;
  }
  prev = flatten(root.right, prev);
  prev = flatten(root.left, prev);
  root.right = prev;
  root.left = null;
  prev = root;
  return prev;
}

Time: $O(N)$
Space: $O(h)$

Iteration (Preorder)

Use a stack to simulate the preorder traversal, but we also use peek operation to set current node’s right child to its left child.

1
2
3
4
5
    1
   / \
  2   5
 / \   \
3   4   6

Note:

  • Before peeking, remember to check if the stack is empty.
  • Actually, we don’t have to check the stack’s size before calling peek(). If the stack is empty, it will just return null. But it does not hold true for Stack‘s peek().
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
public void flatten(TreeNode root) {
  if (root == null) {
    return;
  }
  Stack<TreeNode> stack = new Stack<>();
  stack.push(root);
  while (stack.size() > 0) {
    TreeNode p = stack.pop();
    
    if (p.right != null) { // push right first
      stack.push(p.right);
    }
    
    if (p.left != null) { // push left then
      stack.push(p.left);
    }
    
    if (stack.size() > 0) { // peek the left
      p.right = stack.peek();
    }
    
    p.left = null; // Remember to set left as null
  }
}

Time: $O(N)$
Space: $O(h)$