LeetCode – Refresh – Populating Next Right Pointers in Each Node I and II

 1 /**
 2  * Definition for binary tree with next pointer.
 3  * struct TreeLinkNode {
 4  *  int val;
 5  *  TreeLinkNode *left, *right, *next;
 6  *  TreeLinkNode(int x) : val(x), left(NULL), right(NULL), next(NULL) {}
 7  * };
 8  */
 9 class Solution {
10 public:
11     TreeLinkNode *getNext(TreeLinkNode *root) {
12         if (!root) return NULL;
13         if (root->left) return root->left;
14         if (root->right) return root->right;
15         return getNext(root->next);
16     }
17     void connect(TreeLinkNode *root) {
18         if (!root) return;
19         TreeLinkNode *current = root;
20         while (current) {
21             root = current;
22             while (root) {
23                 if (root->left && root->right) {
24                     root->left->next = root->right;
25                     root->right->next = getNext(root->next);
26                 } else if (root->left || root->right) {
27                     getNext(root)->next = getNext(root->next);
28                 }
29                 root = root->next;
30             }
31             current = getNext(current);
32         }
33     }
34 };