Populating Next Right Pointers in Each Node II.cpp

/**
 * 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)return;
        auto curr = root;
        while(curr)
        {
            curr = findNext(curr);
            if(!curr)break;
            TreeLinkNode* head = curr->left? curr->left: curr->right;
            while(curr)
            {
                if(curr->left && curr->right)
                    curr->left->next = curr->right;
                TreeLinkNode* next = findNext(curr->next);
                if(!next)break;
                if(curr->right)curr->right->next = next->left? next->left: next->right;
                else curr->left->next = next->left? next->left: next->right;
                curr = next;
            }
            curr = head;
        }
    }

private:
    TreeLinkNode* findNext(TreeLinkNode* curr)
    {
        while(curr)
        {
            if(curr->left || curr->right)
                return curr;
            curr = curr->next;
        }
        return curr;
    }
};