Populating Next Right Pointers in Each Node II

 

Follow up for problem "Populating Next Right Pointers in Each Node".

What if the given tree could be any binary tree? Would your previous solution still work?

Note:

• You may only use constant extra space.

For example,
Given the following binary tree,

         1
       /  \
      2    3
     / \    \
    4   5    7

After calling your function, the tree should look like:

         1 -> NULL
       /  \
      2 -> 3 -> NULL
     / \    \
    4-> 5 -> 7 -> NULL

 

Obviously, the thought of this question is also disturbed by the former one. First, post the wrong code:

 

 

/**
 * 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) {
        // Note: The Solution object is instantiated only once and is reused by each test case.
        TreeLinkNode *tmp;
        if(root==NULL)   return;
        if(root->left!=NULL){
            if(root->right!=NULL)
                root->left->next=root->right;
            else{
                tmp=root->next;
                while(tmp!=NULL){
                    if(tmp->left!=NULL){
                        root->left->next=tmp->left;
                        break;
                    }
                    if(tmp->right!=NULL){
                        root->left->next=tmp->right;
                        break;
                    }
                    tmp=tmp->next;
                }       
            }
        }
        
        if(root->right!=NULL){
            tmp=root->next;
            while(tmp!=NULL){
                if(tmp->left!=NULL){
                    root->right->next=tmp->left;
                    break;
                }
                if(tmp->right!=NULL){
                    root->right->next=tmp->right;
                    break;
                }
                tmp=tmp->next;
            }    
        }  
        connect(root->left);
        connect(root->right);
    }
};

 

 

 

Input:	{2,1,3,0,7,9,1,2,#,1,0,#,#,8,8,#,#,#,#,7}
Output:	{2,#,1,3,#,0,7,9,1,#,2,1,0,#,7,#}
Expected:	{2,#,1,3,#,0,7,9,1,#,2,1,0,8,8,#,7,#}

Take notice that:

 

        connect(root->left);
        connect(root->right);

 

The traversal is DFS, therefore, the right part of some layer has been linked. For example, 0->7->9....1 has formed when visiting 7:1->0. However, the next hop of 0 couldn't be linked because 9-->1 hasn't been connected. So the  the right code is:

 

/**
 * 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) {
        // Note: The Solution object is instantiated only once and is reused by each test case.
        TreeLinkNode *tmp,rt,curnode;
        if(root==NULL)   return;
        rt=root;
        while(rt!=NULL){
            if(rt->left!=NULL){
                curnode=rt->left;
                break;
            }
            if(rt->right!=NULL){
                curnode=rt->right;
                break;
            }
            rt=rt->next;
        }
        
        
        if(root->left!=NULL){
            if(root->right!=NULL)
                root->left->next=root->right;
            else{
                tmp=root->next;
                while(tmp!=NULL){
                    if(tmp->left!=NULL){
                        root->left->next=tmp->left;
                        break;
                    }
                    if(tmp->right!=NULL){
                        root->left->next=tmp->right;
                        break;
                    }
                    tmp=tmp->next;
                }       
            }
        }
        
        if(root->right!=NULL){
            tmp=root->next;
            while(tmp!=NULL){
                if(tmp->left!=NULL){
                    root->right->next=tmp->left;
                    break;
                }
                if(tmp->right!=NULL){
                    root->right->next=tmp->right;
                    break;
                }
                tmp=tmp->next;
            }    
        }  
        
        connect(root->right);
        connect(root->left);
    }
};

A concise code could be written as:

 

 

/**
 * 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) {
        // Note: The Solution object is instantiated only once and is reused by each test case.
        TreeLinkNode *rt=NULL,*nextnode=NULL;
        if(root==NULL)   return;
        rt=root->next;
        while(rt!=NULL){ //Find the next node for the rightest child of all children
            if(rt->left!=NULL){
                nextnode=rt->left;
                break;
            }
            if(rt->right!=NULL){
                nextnode=rt->right;
                break;
            }
            rt=rt->next;
        }
        if(root->right!=NULL){
            root->right->next=nextnode;
            if(root->left!=NULL)
                root->left->next=root->right;
        }
        else if(root->left!=NULL)
            root->left->next=nextnode;

        connect(root->right);
        connect(root->left);
    }
};

Python Version:

"""
# Definition for a Node.
class Node:
    def __init__(self, val, left, right, next):
        self.val = val
        self.left = left
        self.right = right
        self.next = next
"""
class Solution:
    def connect(self, root):
        if (root == None):
            return root
        layers=[[root],[]]
        cur,nxt = 0,1 #range in 0 and 1
        while (layers[cur]):
            layersLen = len(layers[cur])
            for i in range(layersLen):
                layers[cur][i].next = layers[cur][i+1] if i+1<layersLen else None
                if (layers[cur][i].left):
                    layers[nxt].append(layers[cur][i].left)
                if (layers[cur][i].right):
                    layers[nxt].append(layers[cur][i].right)
            layers[cur].clear()
            cur,nxt=nxt,cur
        return root