Populating Next Right Pointers in Each Node II

------QUESTION------

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------

class Solution {
public:
    void connect(TreeLinkNode *root) {
        TreeLinkNode* parent;
        TreeLinkNode* current;
        TreeLinkNode* nextParent = root;
   
        while(1){
           parent = nextParent;
           nextParent = NULL;
           while(parent){ //find nextParent
               if(parent->left) {
                   nextParent = parent->left;
                   break;
               }
               if(parent->right){
                   nextParent = parent->right;
                   break;
               }
               parent = parent->next;
           } 
           current = nextParent;
           if(!current) return;
   
           while(parent){
               if(current == parent->left){//add next pointer of left child
                   if(parent->right) current->next = parent->right;
                   else{ //find next until parent equals null
                       parent = parent->next;
                       while(parent){ 
                           if(parent->left) {
                               current->next = parent->left;
                               break;
                           }
                           if(parent->right){
                               current->next = parent->right;
                               break;
                           }
                           parent = parent->next;
                       }
                   }
               }
               else{ //add next pointer of right child
                   parent = parent->next;
                   while(parent){ //find next until parent equals null
                       if(parent->left) {
                           current->next = parent->left;
                           break;
                       }
                       if(parent->right){
                           current->next = parent->right;
                           break;
                       }
                       parent = parent->next;
                   }
               }
               current = current->next;
           } // end of level
        }
    }
};