Description

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,

tree

After calling your function, the tree should look like:

linked tree

Solution

Iterative, time O(n), space O(1)

在题1的基础上，稍微变种即可。

/**
 * 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) {
        if (root == null) {
            return;
        }
        
        TreeLinkNode curr = root;
        
        while (curr != null) {
            TreeLinkNode nextLevelStart = null;
            TreeLinkNode nextLevelPrev = null;
            
            while (curr != null) {
                if (curr.left != null) {
                    if (nextLevelPrev != null) {
                        nextLevelPrev.next = curr.left;
                    } else {
                        nextLevelStart = curr.left;
                    }    
                    
                    nextLevelPrev = curr.left;
                }
                
                if (curr.right != null) {
                    if (nextLevelPrev != null) {
                        nextLevelPrev.next = curr.right;
                    } else {
                        nextLevelStart = curr.right;
                    }
                    
                    nextLevelPrev = curr.right;
                }
                // move to the next node
                curr = curr.next;
            }
            // move to the next level
            curr = nextLevelStart;
        }
    }
}