【leetcode】Populating Next Right Pointers in Each Node 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     void connect(TreeLinkNode *root) {
12        
13         if(root==NULL)
14         {
15             return;
16         }
17        
18         if(root->left!=NULL)
19         {
20             if(root->right!=NULL)
21             {
22                 root->left->next=root->right;
23             }
24             else
25             {
26                 TreeLinkNode *tmp=root->next;
27                 root->left->next=NULL;
28                 while(tmp!=NULL)
29                 {
30                     if(tmp->left!=NULL)
31                     {
32                         root->left->next=tmp->left;
33                         break;
34                     }
35                    
36                     if(tmp->right!=NULL)
37                     {
38                         root->left->next=tmp->right;
39                         break;
40                     }
41                     tmp=tmp->next;
42                 }
43             }
44         }
45        
46         if(root->right!=NULL)
47         {
48             TreeLinkNode *tmp=root->next;
49             root->right->next=NULL;
50             while(tmp!=NULL)
51             {
52                 if(tmp->left!=NULL)
53                 {
54                     root->right->next=tmp->left;
55                     break;
56                 }
57                    
58                 if(tmp->right!=NULL)
59                 {
60                     root->right->next=tmp->right;
61                     break;
62                 }
63                 tmp=tmp->next;
64             }
65         }
66        
67        
68         connect(root->right);
69         connect(root->left);
70        
71     }
72 };