117. Populating Next Right Pointers in Each Node II

/**
 * 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) {
        TreeLinkNode* head = root;
        while (head) {
            TreeLinkNode* nextHead = NULL, *pre = NULL;
            while (head) {
                TreeLinkNode *left = head->left, *right = head->right;
                if (left) {
                    if (nextHead == NULL)
                        nextHead = left;
                    else
                        pre->next = left;
                    pre = left;
                }
                if (right) {
                    if (nextHead == NULL)
                        nextHead = right;
                    else
                        pre->next = right;
                    pre = right;
                }
                head = head->next;
            }
            head = nextHead;
        }
    }
};