LeetCode: 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

---

We can use a tracker to record the starting node of each level and then linking nodes in each level.

C++ Code:

/*
 * func: populate_right
 * goal: populate each next pointer to its next right node
 * @param root: root node of the tree
 * return:
 */
/*
 * Use pointer to record the prev node in the process of linking in current level
 * and a next node to record starting node in the next level
 * complexity: time O(n), space O(1)
 */
void populate_right(TreeLinkNode *root){
    TreeLinkNode *iter = root;
    while(iter){
        TreeLinkNode *prev_node = nullptr;
        TreeLinkNode *next_node = nullptr;
        while(iter){
            if(next_node == nullptr){
                if(iter->left)
                    next_node = iter->left;
                else
                    next_node = iter->right;
            }
            
            if(iter->left){
                if(prev_node)
                    prev_node->next = iter->left;
                prev_node = iter->left;
            }
            
            if(iter->right){
                if(prev_node)
                    prev_node->next = iter->right;
                prev_node = iter->right;
            }
            
            iter = iter->next;
        }
        iter = next_node;
    }
}

Python Code:

# func: populate each next pointer to its next right node
# @param root: root node of the tree
# @return:
def populate_right(root):
    node = root
    while node:
        prev_node = None
        next_node = None
        while node:
            #Setting the starting node for the next level
            if next_node is None:
                if node.left:
                    next_node = node.left
                else:
                    next_node = node.right
            #Link the node in the current level
            if node.left:
                if prev_node:
                    prev_node.next = node.left
                prev_node = node.left

            if node.right:
                if prev_node:
                    prev_node.next = node.right
                prev_node = node.right

            node = node.next
        node = next_node

Reference Link:

1. Connect nodes at same level using constant extra space | GeeksforGeekshttp://www.geeksforgeeks.org/connect-nodes-at-same-level-with-o1-extra-space/