【leetcode刷题笔记】Populating Next Right Pointers in Each Node II

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

---

最开始的想法是用一个队列和一个列表。当前遍历的层的元素放在队列中，在遍历当前层的时候，把下一层元素的next指针设置好，并且把下一层的元素放在列表中。之所以要用列表，是因为当遍历一个node的孩子的时候，它在next指针下的前驱就是列表中的最后一个元素，如果用队列的话就取不出最后一个元素，所以这里要用列表。在一层遍历结束后，要把列表中的元素全部放入到队列中，开始下一层的遍历，代码如下：

/**
 * 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;
        
        Queue<TreeLinkNode> currLevel = new LinkedList<TreeLinkNode>();
        List<TreeLinkNode> nextLevel = new ArrayList<TreeLinkNode>();
        
        currLevel.add(root);
        
        while(!currLevel.isEmpty()){
            while(!currLevel.isEmpty()){
                TreeLinkNode temp = currLevel.poll();
                if(temp.left != null){
                    if(!nextLevel.isEmpty() && nextLevel.get(nextLevel.size()-1).next == null)
                        nextLevel.get(nextLevel.size()-1).next = temp.left;
                    nextLevel.add(temp.left);
                }
                if(temp.right != null){
                    if(!nextLevel.isEmpty() && nextLevel.get(nextLevel.size()-1).next == null)
                        nextLevel.get(nextLevel.size()-1).next = temp.right;
                    nextLevel.add(temp.right);
                }
            }
            List<TreeLinkNode> passby = new ArrayList<TreeLinkNode>(nextLevel);
            currLevel.addAll(passby);
            nextLevel.clear();
        }        
    }
}

但是题目中要求要常数的额外空间，上述代码居然也AC了，可见leetcode要求不是特别严格。

那么我们怎么把上述代码修改成原地算法呢？

首先，对于上述用到的列表，其实我们每次只需要列表的最后一个元素，它是当前遍历节点的孩子的前驱。如题目中所示的例子，在遍历2的右子5的时候，只要知道它的前驱是4，即刚刚遍历过的2的左子即可；同理，在遍历节点3的时候，要知道它的右孩子的前驱，只要知道刚刚遍历到的下一层节点5就可以了；所以我们只需要一个TreeNode prev记录刚刚通过父节点遍历到的下一层的节点就可以了。

其次，对于上述用到的队列，队列中元素的先后顺序其实我们在遍历上一层的时候确定这一层元素的next指针，那么在遍历这一层的时候就可以利用这个next指针了，但是仍然需要知道这一层的起点指针，这个起点指针也可以在遍历上一层的时候确定。

修改后的代码如下:

/**
 * 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 first = root;
        while(first != null){
            TreeLinkNode pointer = first;
            first = null;
            TreeLinkNode prev = null;
            
            while(pointer != null){
                //更新first               
                if(pointer.left != null){
                    if(first == null){
                        first = pointer.left;
                        prev = first;
                    }else{
                        prev.next = pointer.left;
                        prev = prev.next;
                    }
                }
                
                if(pointer.right != null){
                    if(first == null){
                        first = pointer.right;
                        prev = first;
                    }else{
                        prev.next = pointer.right;
                        prev = prev.next;
                    }
                }
                
                pointer = pointer.next;
            }
        }
    }
}