Leetcode: Flatten Binary Tree to Linked List

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

For example,
Given

         1
        / \
       2   5
      / \   \
     3   4   6

The flattened tree should look like:

   1
    \
     2
      \
       3
        \
         4
          \
           5
            \
             6

click to show hints.

Hints:

If you notice carefully in the flattened tree, each node's right child points to the next node of a pre-order traversal.

Understand the problem:
The problem asks for flattening a binary tree to a linked list in-place. Given the example shown above, it is actually a pre-order traversal. Note that we are asked to do it in-place meaning we cannot allocate a new linked list to store the result. Furthermore, the binary tree is "right" flatten, meaning each parent points to its right child. 

Solution:
The basic idea is to use pre-order traversal. For a parent node, merge its left subtree to its right child, append the rest of the rest subtree to the right most child of the left substree. Then recursively move to its right child. 

Code (Java):

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void flatten(TreeNode root){
        if (root == null) return;
        
        if (root.left != null) {
            TreeNode rChild = root.right;
            root.right = root.left;
            root.left = null;
            TreeNode rMost = root.right;
            while (rMost.right != null) {
                rMost = rMost.right;
            }
            rMost.right = rChild;
        }
        
        flatten(root.right);
    }
}

Iterative Solution:
Since preorder traversal can be done iteratively, can we solve this problem in an iterative way? The answer is yes. In similar to the preorder traversal, we can employ a stack to solve the problem. 

The basic idea is for a node, if it has right child, push it to the stack first. Then check its left child, if not null, connect it to the right end; else, pop a node from the stack and append it to the right end.

Code (Java):

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void flatten(TreeNode root){
        if (root == null) return;
        Stack<TreeNode> stack = new Stack<TreeNode>();
        
        TreeNode p = root;
        
        while (p != null || !stack.empty()) {
            if (p.right != null) {
                stack.push(p.right);
            }
            
            if (p.left != null) {
                p.right = p.left;
                p.left = null;
            } else if (!stack.empty()) {
                TreeNode temp = stack.pop();
                p.right = temp;
            }
            
            p = p.right;
        }
    }
}

A better solution:
Here is a better solution without using stacks. So the space complexity is O(1). And the code structure is almost the same as recursive solution.

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void flatten(TreeNode root){
        if (root == null) return;
        TreeNode p = root;
        
        while (p.left != null || p.right != null) {
            if (p.left != null) {
                TreeNode rChild = p.right;
                p.right = p.left;
                p.left = null;
                TreeNode rMost = p.right;
                while (rMost.right != null) {
                    rMost = rMost.right;
                }
                rMost.right = rChild;
            }
            
            p = p.right;
        }
    }
}

Summary:
In this post, we learned how to flatten a tree to a linked list. Note the idea of pre-order traversal, and its variants. Also be careful that when you move a node in a tree, all its subtrees will be moved as well, similar to linked list. 

Update on 4/15/15:
An iterative solution which is just transformed from the recursive one. 

Code (Java):

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void flatten(TreeNode root) {
        if (root == null) {
            return;
        }        
        
        Stack<TreeNode> stack = new Stack<TreeNode>();
        stack.push(root);
        
        while (!stack.isEmpty()) {
            TreeNode curr = stack.pop();
            
            if (curr.right != null) {
                stack.push(curr.right);
            }
            
            if (curr.left != null) {
                stack.push(curr.left);
            }
            
            if (curr.left != null) {
                TreeNode rChild = curr.right;
                curr.right = curr.left;
                curr.left = null;
                
                // find the right most child
                TreeNode rightMostChild = curr.right;
                while (rightMostChild.right != null) {
                    rightMostChild = rightMostChild.right;
                }
                rightMostChild.right = rChild;
            }
        }
    }
}

However, we could see from this solution that it is not very inefficient; it will repeatedly add nodes into the stack and pop them out.