117. Populating Next Right Pointers in Each Node II Leetcode Java

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

Solution:

It is the general problem of the previous problem 116. Populating Next Right Pointers in Each Node.  

Here the Tree may be not a perfect binary tree. So when we finished processing the current level, we have to jump to the firstNode of the next level. Here I use firstNode to track the next level's firstNode, the idea is during the traversal of current level, if there is child for any node, check if next level's firstNode has been assigned(firstNode ? =null) if not assigned(firstNode==null) assign this child as the firstNode of next level. Another tricky part is that, we need to maintain a pre variable in order to connect the next level's nodes. For example in the problem, in order to connect 5 and 7, unlike the perfect BT(cur.left.next->cur.right, cur.right.next-> cur.next.left), there is no pattern to follow, so I use a pre variable to store the pre node. If cur.next==null, which means the current level has been finished, the iterator will jump to next level's firstNode.

  public void connect(TreeLinkNode root) {  
    if(root==null) return;  
    TreeLinkNode cur=root;  
    TreeLinkNode pre=null;  
    TreeLinkNode firstNode=null;  
    while(cur!=null){  
      if(cur.left!=null){  
        if(pre!=null) pre.next=cur.left;  
        pre=cur.left;  
        if(firstNode==null) firstNode=cur.left;  
      }  
      if(cur.right!=null){  
        if(pre!=null) pre.next=cur.right;  
        pre=cur.right;  
        if(firstNode==null) firstNode=cur.right;  
      }  
      if(cur.next==null){  
        pre=null;  
        cur=firstNode;  
        firstNode=null;  
      }  
      else cur=cur.next;  
    }  
   }