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

59 阅读3分钟

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

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:
            nodes = []
            nextLayer = []
            for node in queue:
                nodes.append(node)
                if node.left and node.right:
                    nextLayer.append(node.left)
                    nextLayer.append(node.right)

            for i in range(0, len(nodes) - 1):
                nodes[i].next = nodes[i + 1]

            nodes[-1].next = None
            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:
            # 遍历每一层
            node = cur
            while node:
                # 当前节点的左孩子指向右孩子
                node.left.next = node.right
                # 当前节点的右孩子指向当前节点next的左孩子
                if node.next:
                    node.right.next = node.next.left
                
                # 继续往右
                node = node.next
            # 下一层
            cur = cur.left
        
        return root

对于每一层给定的两个节点node1node2来说,所做的事情都是相同的:

  • 首先,node1.next = node2
  • 然后继续分别操作node1node2的左右孩子之间的next指向
  • 此外,还需要处理node1.right.next = node2.left

递归代码如下:

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

        def helper(node1, node2):
            if not node1 or not node2:return None

            node1.next = node2

            helper(node1.left, node1.right)
            helper(node2.left, node2.right)
            helper(node1.right, node2.left)
        
        helper(root.left, root.right)

        return root