116. Populating Next Right Pointers in Each Node

Leetcode Tree Depth-first Search

Given a binary tree

public class TreeLinkNode {
  int val;
  TreeLinkNode left, right, next;
  TreeLinkNode(int x) { val = x; }
}

Populate each next pointer to point to its next right node. If there is no next right node, the next pointer should be set to NULL.

Initially, all next pointers are set to NULL.

Note:

• You may only use constant extra space.
• Recursive approach is fine, implicit stack space does not count as extra space for this problem.
• You may assume that it is a perfect binary tree (ie, all leaves are at the same level, and every parent has two children).

Example:

Given the following perfect binary tree,

     1
   /  \
  2    3
 / \  / \
4  5  6  7

After calling your function, the tree should look like:

     1 -> NULL
   /  \
  2 -> 3 -> NULL
 / \  / \
4->5->6->7 -> NULL

分析

这道题目是二叉树层序遍历(102. Binary Tree Level Order Traversal)的变形。在层序遍历的同时，要求每个节点的next指向右边的节点。既然二叉树的层序遍历有递归和迭代两种形式，这里也写上这两种形式。

迭代的形式，变量prev保存前一个遍历的节点， prev.next指向当前节点。

public void connect(TreeLinkNode root) {
    if (root == null) return;
    Queue<TreeLinkNode> queue = new LinkedList<>();
    queue.offer(root);
    TreeLinkNode cur, prev;
    while(!queue.isEmpty()) {
        int size = queue.size();
        prev = null;
        for (int i = 0; i < size; i++) {
            cur = queue.poll();
            if (prev != null)
                prev.next = cur;
            prev = cur;
            // push nodes of the next level
            if (cur.left != null) queue.offer(cur.left);
            if (cur.right != null) queue.offer(cur.right);
        }
    }
}

递归的形式，维持一个变量prev为每一层的前一个遍历的节点。

private List<TreeLinkNode> prev;
public void connect(TreeLinkNode root) {
    prev = new ArrayList<>();
    levelOrderTraversal(root, 0);
}

private void levelOrderTraversal(TreeLinkNode root, int height) {
    if (root == null) return;
    if (height >= prev.size()) prev.add(root);
    else prev.get(height).next = root;
    prev.set(height, root);
    levelOrderTraversal(root.left, height + 1);
    levelOrderTraversal(root.right, height + 1);
}

最快度的方案也是最优美的方案利用了题目中的完全二叉树的提示。到目前为止，还没有利用上这一条约束。加上这一条约束以后，其实可以从左到右，从上到下依次直接连接节点：

public void connect(TreeLinkNode root) {
    if (root == null) return;
    while (root != null) {
        TreeLinkNode firstNode = root;
        while(root != null){
            // 左子节点指向右子节点
            if(root.left != null) 
                root.left.next = root.right;
            // 右子节点指向下一个(右边)节点的左节点
            if(root.right != null && root.next != null) 
                root.right.next = root.next.left;
            // 往右移动
            root = root.next;
        }
        // 移动到下一层
        root = firstNode.left;
    }
}