LeetCode/Scripts/117.cpp at master · zzz0906/LeetCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
public:
    Node* connect(Node* root) {
        if (!root) return NULL;
        if (root->left){
            if (root->right) {
                root->left->next = root->right;
                Node* tmp = root->next;
                while (tmp != NULL && (tmp->left == NULL || tmp->right == NULL)) tmp = tmp->next;
                if (tmp){
                    if (tmp->left == NULL)
                        root->right->next = tmp->right;
                    else
                        root->right->next = tmp->left;
                }
                else root->right->next = NULL;
            }else{
                Node* tmp = root->next;
                while (tmp != NULL && (tmp->left == NULL || tmp->right == NULL)) tmp = tmp->next;
                if (tmp){
                    if (tmp->left == NULL)
                        root->left->next = tmp->right;
                    else
                        root->left->next = tmp->left;
                }
                else root->left->next = NULL;
            }
        } if (root->right){
            Node* tmp = root->next;
            while (tmp != NULL && (tmp->left == NULL || tmp->right == NULL)) tmp = tmp->next;
            if (tmp){
                if (tmp->left == NULL)
                    root->right->next = tmp->right;
                else
                    root->right->next = tmp->left;
            }
            else root->right->next = NULL;
        }
        connect(root->left);
        connect(root->right);
        return root;
    }
};