Flatten Binary Tree to Linked List @LeetCode

package Level4;

import Utility.TreeNode;

/**
 * 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.
 */
public class S114 {

	public static void main(String[] args) {
		TreeNode r1 = new TreeNode(1);
		TreeNode r2 = new TreeNode(2);
		TreeNode r3 = new TreeNode(5);
		TreeNode r4 = new TreeNode(3);
		TreeNode r5 = new TreeNode(4);
		TreeNode r6 = new TreeNode(6);
		r1.left = r2;
		r1.right = r3;
		r2.left = r4;
		r2.right = r5;
		r3.right = r6;
		
		flatten(r1);
		while(r1 != null){
			System.out.print(r1.val + " ");
			r1 = r1.right;
		}
	}

	public static void flatten(TreeNode root) {
        rec(root);
    }
	
	/**
	             1
		        / \
		       2   5
		      / \   \
		     3   4   6
	 */
	public static TreeNode rec(TreeNode root){
		if(root==null){
			return root;
		}
		
		TreeNode leftSub = rec(root.left);		// leftSub为左子树的根节点，比如现在指向2节点，这颗子树是  2->





3->4
		TreeNode rightSub = rec(root.right);	// rightSub为右子树的根节点，比如现在指向5节点，这颗子树是 5->6
		
		root.right = leftSub;		// 使得 1 -> 2->3->4
		root.left = null;
		
		TreeNode rightMost = leftSub;		// 现在要连接4和rightSub，即5。因此需要找到4，即最右端的节点
		
		if(rightMost != null){
			while(rightMost.right != null){	// 寻找最右端节点
				rightMost = rightMost.right;
			}
		}
		
		if(rightMost != null){		// 存在最右端节点，就连接4和rightSub(5)
			rightMost.right = rightSub;
		}else{							// 如果不存在就用根节点直接连接rightSub(5)
			root.right = rightSub;
		}
		
		return root;
	}
}