Leetcode: Populating Next Right Pointers in Each Node II

Follow up for problem "Populating Next Right Pointers in Each Node".

What if the given tree could be any binary tree? Would your previous solution still work?

Note:

• You may only use constant extra space.

For example,
Given the following binary tree,

         1
       /  \
      2    3
     / \    \
    4   5    7

After calling your function, the tree should look like:

         1 -> NULL
       /  \
      2 -> 3 -> NULL
     / \    \
    4-> 5 -> 7 -> NULL

Understand the problem:
This is a follow-up problem of the problem 1. The mainly difference is now the tree is not perfect binary tree any more. Note that you are still required to use constant extra space. 
So if space is not a problem, we may still use a queue to do a BFS. 

Solution:
Since now each node may not have its left or right child, the previous approach may not work. The key of the solution is to find a valid next node. 

Code (Java):

/**
 * Definition for binary tree with next pointer.
 * public class TreeLinkNode {
 *     int val;
 *     TreeLinkNode left, right, next;
 *     TreeLinkNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void connect(TreeLinkNode root) {
        if (root == null) {
            return;
        }
        
        if (root.left != null && root.right != null) {
            root.left.next = root.right;
        }
        
        TreeLinkNode p = root;
        
        while (p.next != null && p.next.left == null && p.next.right == null) {
            p = p.next;
        }
        
        
        if (root.right != null && p.next != null) {
            if (p.next.left != null) {
                root.right.next = p.next.left;
            } else if (p.next.right != null){
                root.right.next = p.next.right;
            }
        } else if (root.left != null && p.next != null) {
            if (p.next.left != null) {
                root.left.next = p.next.left;
            } else if (p.next.right != null) {
                root.left.next = p.next.right;
            }
        }
        
        connect(root.right);
        connect(root.left);
    }
}

Discussion:
The only tricks here is to go to its right child first then its left right, i.e, connect right pointers at each level backward. That is because we wanna find the first valid node. We can use an example to illustrate this. If we go to the left child before right one, 

Input:	{2,1,3,0,7,9,1,2,#,1,0,#,#,8,8,#,#,#,#,7}
Output:	{2,#,1,3,#,0,7,9,1,#,2,1,0,#,7,#}
Expected:	{2,#,1,3,#,0,7,9,1,#,2,1,0,8,8,#,7,#}

Why? That is because when we are at Node 7, we figure out the p pointer points to node 9. However, since now 9 is not connected to 1 yet, so node 7's right child, 0 cannot point to node 1's left child, which is 8. That is why we must traverse the right node first to make 9 connects to 1 before we check the node 7.


Update on 9/18/15:

/**
 * Definition for binary tree with next pointer.
 * public class TreeLinkNode {
 *     int val;
 *     TreeLinkNode left, right, next;
 *     TreeLinkNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void connect(TreeLinkNode root) {
        if (root == null) {
            return;
        }
        
        if (root.left != null && root.right != null) {
            root.left.next = root.right;
        }
        
        if (root.next != null) {
            TreeLinkNode p = root;
            while (p.next != null && p.next.left == null && p.next.right == null) {
                p = p.next;
            }
            
            if (p.next != null) {
                p = p.next;
            }
            
            if (root.right != null) {
                if (p.left != null) {
                    root.right.next = p.left;
                } else if (p.right != null) {
                    root.right.next = p.right;
                }
            } else if (root.left != null) {
                if (p.left != null) {
                    root.left.next = p.left;
                } else if (p.right != null) {
                    root.left.next = p.right;
                }
            }
        }
        
        connect(root.right);
        connect(root.left);
    }
}

Update on 2/8/16:
Constant space solution:

/**
 * Definition for binary tree with next pointer.
 * public class TreeLinkNode {
 *     int val;
 *     TreeLinkNode left, right, next;
 *     TreeLinkNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void connect(TreeLinkNode root) {
        if (root == null) {
            return;
        }
        
        TreeLinkNode curr = root;
        TreeLinkNode dummyNode = new TreeLinkNode(0);
        TreeLinkNode pre = dummyNode;
        
        while (curr != null) {
            TreeLinkNode p = curr;
            while (p != null) {
                if (p.left != null) {
                    pre.next = p.left;
                    pre = pre.next;
                }
                
                if (p.right != null) {
                    pre.next = p.right;
                    pre = pre.next;
                }
                
                p = p.next;
                
                if (p == null) {
                    curr = dummyNode.next;
                    pre = dummyNode;
                    dummyNode.next = null;
                }
            }
        }
        
    }
}