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.

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;
           
        flattenRecur(root);
        TreeNode end = root;
        TreeNode cur = root.left;
        
        while (cur!=null){
            end.right = cur;
            end.left = null;
            end = cur;
            cur = cur.left;
        }
    }
    
    public TreeNode flattenRecur(TreeNode curNode){
        TreeNode leftEnd = null;
        TreeNode rightEnd = null;
        
        //If curNode is a leaf node, then return
        if (curNode.left==null && curNode.right==null){
            return curNode;
        }
        
        //If curNode is not a leaf node, then visit its left child first.
        if (curNode.left!=null)
            leftEnd = flattenRecur(curNode.left);
        
        //If curNode has right child, then move the right child tree to the end of ths list.
        if (curNode.right!=null)
            rightEnd = flattenRecur(curNode.right);
        
        if (leftEnd==null){
            curNode.left = curNode.right;
            curNode.right = null;
            return rightEnd;
        } else {
            if (rightEnd!=null){
                leftEnd.left = curNode.right;
                curNode.right = null;
                return rightEnd;
            } else
                return leftEnd;
        }
    }
}

For a node, first flatten its left tree, then flaten its right tree. Get the end of the flated list of both of its left and right tree, then put the list of its right child to the end of the list of its left child. We then flaten the tree with current node (root). We return the end of the flatted list to the upper level.

NOTE: We need to be consider about the situations that the left or right child tree is NULL.