Populating Next Right Pointers in Each Node II | Leetcode

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

---

Solution: The trick is to use the linked list of a level to fill up the next lower level. Maintain the first node (curFirst) of the linked list. At the beginning of each loop, set the current node (curNode) to the first node. Then for each left and right children, link them up as you traverse the level from left to right. For the first non-null child node you find, set it to be the first node.

/**
 * Definition for binary tree with next pointer.
 * public class TreeLinkNode {
 *     int val;
 *     TreeLinkNode left, right, next;
 *     TreeLinkNode(int x) { val = x; }
 * }
 */
public class Solution {
    public void connect(TreeLinkNode root) {
        TreeLinkNode curFirst = root;
        TreeLinkNode curNode;
        TreeLinkNode nextNode;
        
        while (curFirst != null) {
            curNode = curFirst;
            curFirst = null;
            nextNode = null;
            
            do {
                if (curNode.left != null) {
                    if (nextNode == null) {
                        nextNode = curNode.left;
                        if (curFirst == null) curFirst = nextNode;
                    } else {
                        nextNode = nextNode.next = curNode.left;
                    }
                }
                
                if (curNode.right != null) {
                    if (nextNode == null) {
                        nextNode = curNode.right;
                        if (curFirst == null) curFirst = nextNode;
                    } else {
                        nextNode = nextNode.next = curNode.right;
                    }
                }
                
                curNode = curNode.next;
            } while (curNode != null);
        }
    }
}