[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
问题描述：给定一个二叉树，使每个节点的next指针指向它的右边的节点，和之前的Populating Next Right Pointers in Each Node类似，只是这次的二叉树不是完全二叉树，但是方法和思想与之前的一样。
用一个指针遍历每层的最左节点，用另一个指针遍历前面的指针所在层的所有节点，由于不是完全二叉树，这里引入了另外一个函数find_has_child()，它返回从参数节点开始，该层第一个有孩子的节点。

class Solution {
public:
TreeLinkNode *find_has_child(TreeLinkNode *parent)
{
TreeLinkNode *p = parent;

while(p && !(p->left) && !(p->right))
p = p->next;

return p;
}

void connect(TreeLinkNode *root) {
// IMPORTANT: Please reset any member data you declared, as
// the same Solution instance will be reused for each test case.
TreeLinkNode *parent = root;

if(!parent)
return;
parent->next = NULL;

while(parent) {
TreeLinkNode *level_parent = parent;
while(level_parent->next) {
TreeLinkNode *lchild = level_parent->left;
TreeLinkNode *rchild = level_parent->right;
if(lchild && !rchild) {
TreeLinkNode *p = find_has_child(level_parent->next);
if(p) {
if(p->left) {
lchild->next = p->left;
}
else {
lchild->next = p->right;
}
level_parent = p;
continue;
}
else {
lchild->next = NULL;
break;
}
}

if(rchild) {
if(lchild)
lchild->next = rchild;
TreeLinkNode *p = find_has_child(level_parent->next);
if(p) {
if(p->left) {
rchild->next = p->left;
}
else {
rchild->next = p->right;
}
level_parent = p;
continue;
}
else {
rchild->next = NULL;
break;
}
}
level_parent = level_parent->next;
}
if(level_parent->left)
level_parent->left->next = level_parent->right;
else if(level_parent->right)
level_parent->right->next = NULL;

if(parent->left) {
parent = parent->left;
}
else if(parent->right) {
parent = parent->right;
}
else {
TreeLinkNode *p = find_has_child(parent->next);
if(p) {
if(p->left) {
parent = p->left;
}
else if(p->right) {
parent = p->right;
}
}
else {
parent = NULL;
}
}
}
}
};

虽然代码很长，但是逻辑还是很清楚的。
当遍历到一层时，如果有左孩子，没有右孩子，那么左孩子的next就指向父节点的下一个有孩子节点的第一个孩子；如果既有左孩子，又有右孩子，那么左孩子的next指向有孩子，右孩子的next指向父节点的下一个有孩子节点的第一个孩子；如果有右孩子，没有左孩子，那么右孩子的next指向父节点的下一个有孩子节点的第一个孩子，这种情况和第二种情况有重叠；当然，如果没有孩子节点，那么就遍历到下一个节点。
不过上面的代码有些冗余，在两个if中代码相似度很高，可以将上面的两个if简化：

if(lchild && rchild)
lchild->next = rchild;
if(lchild || rchild) {
TreeLinkNode *p = find_has_child(level_parent->next);
TreeLinkNode *q = NULL;
if(lchild && !rchild)
q = lchild;
if(rchild)
q = rchild;
if(p) {
if(p->left)
q->next = p->left;
else
q->next = p->right;
level_parent = p;
continue;
}
else {
q->next = NULL;
break;
}
}