给定一个完美二叉树,其所有叶子节点都在同一层,每个父节点都有两个子节点。二叉树定义如下:
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
对于每一层给定的两个节点node1
和node2
来说,所做的事情都是相同的:
- 首先,
node1.next = node2
- 然后继续分别操作
node1
和node2
的左右孩子之间的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