flatten - facebook code lab - flatten the binary tree - pass test cases

flatten2.java

/**
 * Definition for binary tree
 * class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
	public TreeNode flatten(TreeNode a) {
	    if(a==null || a.val == -1)
	        return null; 
	  
	    TreeNode l_subtree = flatten(a.left); // 
	    TreeNode r_subtree = flatten(a.right);
	    
	    TreeNode lastNode = lastNode(l_subtree); 
	        
	    a.left = null; 
	   
	    if(l_subtree != null)  // do not forget, bug001 - runtime error. 
	    {
	        a.right = l_subtree; 
	        lastNode.right = r_subtree; 
	    }
	    
	    return a; 
	}
	
	public TreeNode lastNode(TreeNode a)
	{
	    TreeNode lastNode = a; 
	    while(lastNode !=null && lastNode.val != -1 &&
	          lastNode.right !=null && lastNode.right.val != -1)
	    {
	       
	       lastNode = lastNode.right; 
	    }
	    
	    return lastNode;
	}
}