Populating Next Right Pointers in Each Node

Description

https://leetcode.com/problems/populating-next-right-pointers-in-each-node/description/

Solution

/**
 * 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 == NULL) return;
        TreeLinkNode* cur = root;
        while(cur != NULL) {
            cur = this->linkSameLayer(cur);
        }
    }
    
    TreeLinkNode* linkSameLayer(TreeLinkNode* parent) {
        TreeLinkNode* traverse = parent->left;
        if (traverse == NULL) return NULL;
        
        TreeLinkNode* ret = traverse;
        int count = 0;
        while (parent != NULL) {
            traverse->next = count % 2 ? parent->left : parent->right;
            traverse = traverse->next;
            parent = count % 2 ? parent : parent->next;
            count += 1;
        }
        return ret;
    }
};