117 - Populating Next Right Pointers in Each Node II

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

Hide Tags
Tree Depth-first Search
Hide Similar Problems
(M) Populating Next Right Pointers in Each Node

跟第一题类似，下一层利用上一层已经建立起来的next指针，来串联同一行的所有节点。
不同的是，每个节点可能有0,1,2个子节点，一共四种情况，
所以 while 的条件不能用之前的是否有左子，只能是 p 是否为空，循环内部则按照不同情形来处理。
因为有些节点可能没有子节点，所以，找自己的最后一个子节点的下一个节点时，
需要沿着p->next，一直找到第一个有子节点的节点，取它的第一个子节点才行。
一行处理完了，找下一行的首节点的时候，也一样，要找到第一个有子节点的节点作为下一行处理的起点。

/**
* Definition for binary tree with next pointer.
* struct TreeLinkNode {
*  int val;
*  TreeLinkNode *left, *right, *next;
*  TreeLinkNode(int x) : val(x), left(NULL), right(NULL), next(NULL) {}
* };
*/
class Solution {
public:
    void connect(TreeLinkNode *root) {
        if(root && ((root->left)|| (root->right)))
        {
            TreeLinkNode* levelHead=root;
            TreeLinkNode* p = levelHead;

            root->next = NULL;
            while(p)
            {
                if(p->right)
                {
                    if(p->left)
                        p->left->next = p->right;
                    p->right->next = nextToMyChild(p);
                }   
                else if(p->left)
                    p->left->next = nextToMyChild(p);
                else
                    ;

                if(NULL != p->next)
                       p = p->next;
                else
                {
                    //In current level, find first node who has at least one child
                    while(levelHead && (NULL == levelHead->left) && (NULL == levelHead->right))
                            levelHead = levelHead->next;
                   
                    //the first child of this node will be the head node of next level.
                    if(levelHead)
                        levelHead = (levelHead->left)?levelHead->left:levelHead->right;

                    p = levelHead;
                }
            }
        }
        return;
    }
   
    TreeLinkNode* nextToMyChild(TreeLinkNode* p)
    {
        TreeLinkNode* pp = p->next;

        //keep finding until hit one has at least one child
        while(pp && (NULL == pp->left) && (NULL == pp->right))
            pp = pp->next;

        return pp ? ((pp->left) ? (pp->left) : ((pp->right)?(pp->right):NULL)) : NULL;
    }
};

36 ms.