117. Populating Next Right Pointers In Each NodeII

364 查看

题目:
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

解答:

public class Solution {
    public void connect(TreeLinkNode root) {
        //重点是记录current, head, prev这三个结点
        TreeLinkNode curt = root;
        TreeLinkNode head = null;
        TreeLinkNode prev = null;
        
        while (curt != null) {
            while (curt != null) {
                if (curt.left != null) {
                    if (prev != null) {
                        prev.next = curt.left;
                    } else {
                        head = curt.left;
                    }
                    prev = curt.left;
                }
                if (curt.right != null) {
                    if (prev != null) {
                        prev.next = curt.right;
                    } else {
                        head = curt.right;
                    }
                    prev = curt.right;
                }
                curt = curt.next;
            }
            curt = head;
            head = null;
            prev = null;
        }
    }
}