Leetcode 117 Populating Next Right Pointers in Each Node 2

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

和之前的有不一样的是 最左边的点不确定 然后每个节点的孩子不确定, 所以要判断的条件会多一些

class Solution {
public:
    void connect(TreeLinkNode *root) {
        if (root == NULL)
            return;
        if (root->right == NULL && root ->left == NULL)
            return;
        TreeLinkNode * p = root;
        TreeLinkNode * current = root;
        int flag = 0;
        while (root ->left || root ->right || root->next)
        {
            p = root; 
            flag = 0;
            while(p!=NULL)
            {
                if(p->left == NULL && p->right == NULL)
                {
                    p = p -> next;
                    continue; 
                }
                else 
                {
                    if (flag == 0)
                    {
                        current = p->left ? p->left:p->right; 
                        if (p->left != NULL && p->right != NULL)
                        {
                            current ->next = p ->right ;
                            current = current ->next ; 
                        }
                        flag = 1;
                    }
                    else 
                    {
                        if(p->left && p->right)
                        {
                            current ->next = p->left; 
                            p->left ->next = p -> right;
                            current = p->right; 
                        }
                        else if(p->left!=NULL && p->right==NULL)
                        {
                            current ->next = p->left; 
                            current = p->left; 
                        }
                        else 
                        {
                            current ->next = p->right; 
                            current = p->right; 
                        }
                    }
                }  
                p = p->next;
            }
            if (root -> left)
                root = root ->left; 
            else if (root ->right)
                root = root ->right; 
            else if(root ->next)
                root = root ->next; 
        }
    }
};