116 - 填充每个节点的下一个右侧节点指针 - python

给定一个完美二叉树，其所有叶子节点都在同一层，每个父节点都有两个子节点。二叉树定义如下：

struct Node {
  int val;
  Node *left;
  Node *right;
  Node *next;
}

填充它的每个 next 指针，让这个指针指向其下一个右侧节点。如果找不到下一个右侧节点，则将 next 指针设置为 NULL。

初始状态下，所有 next 指针都被设置为 NULL。

示例：

输入：{"$id":"1","left":{"$id":"2","left":{"$id":"3","left":null,"next":null,"right":null,"val":4},"next":null,"right":{"$id":"4","left":null,"next":null,"right":null,"val":5},"val":2},"next":null,"right":{"$id":"5","left":{"$id":"6","left":null,"next":null,"right":null,"val":6},"next":null,"right":{"$id":"7","left":null,"next":null,"right":null,"val":7},"val":3},"val":1}

输出：{"$id":"1","left":{"$id":"2","left":{"$id":"3","left":null,"next":{"$id":"4","left":null,"next":{"$id":"5","left":null,"next":{"$id":"6","left":null,"next":null,"right":null,"val":7},"right":null,"val":6},"right":null,"val":5},"right":null,"val":4},"next":{"$id":"7","left":{"$ref":"5"},"next":null,"right":{"$ref":"6"},"val":3},"right":{"$ref":"4"},"val":2},"next":null,"right":{"$ref":"7"},"val":1}

解释：给定二叉树如图 A 所示，你的函数应该填充它的每个 next 指针，以指向其下一个右侧节点，如图 B 所示。

---

完美二叉树满足两个条件：

• 所有叶子节点都在同一层
• 所有的父节点都有两个子节点

题目要求为每个节点设置next指针，让其指向可能的下一个右侧节点。从给定的示意图我们可以很容易的看出，next指针的设置只是在层次遍历的基础上多加了一些操作而已。因此，从最基本的层次遍历出发，在每一层的遍历中分别保存所有的父节点到数组res中。设置next指针就相当于让res[i]指向可能的res[i + 1]，如果res[i + 1]不存在则指向None。

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        if root == None:
            return root

        queue = [root]
        while queue:
            res = []
            nextLayer = []
            for node in queue:
            	# 完美二叉树中左右子树要么都不存在，要么都有
                if not node.left and not node.right:
                    res.append(node)
                elif node.left and node.right:
                    res.append(node)
                    nextLayer.append(node.left)
                    nextLayer.append(node.right)

                for i in range(len(res) - 1):
                    res[i].next = res[i + 1]
                res[-1].next = None

            del queue[0]
            queue = nextLayer
        
        return root

使用层次遍历的方法先获取同一层的节点后，再进行next指针的设置需要额外的存储空间。我们再仔细观察一下完美二叉树设置好next指针后的情况

从上图可以看出，父节点的左子节点的next指针只需要将其指向它的右子节点即可，困难的地方在于不同的父节点的右子节点和左子节点间的连接。如果当前节点为node = TreeNode(2)，那么首先有node.left.next = node.right，即4指向5。而5指向6的next指针在该层无法设置，但是我们可以通过node访问到6，然后将5的next指针指向6即可，即node.right.next = node.next.left。

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        
        if not root:
            return root

       cur = root
       while cur.left:
        	# head用于层间移动设置
            head = cur
            while head:
                head.left.next = head.right
                if head.next:
                    head.right.next = head.next.left
                
                head = head.next
            cur= cur.left
        
        return root