填充每个节点的下一个右侧节点指针 II

标签: 深度优先搜索 广度优先搜索 链表 二叉树

难度: Medium

给定一个二叉树:

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

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

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

示例 1:

输入:root = [1,2,3,4,5,null,7]
输出:[1,#,2,3,#,4,5,7,#]
解释:给定二叉树如图 A 所示,你的函数应该填充它的每个 next 指针,以指向其下一个右侧节点,如图 B 所示。序列化输出按层序遍历顺序(由 next 指针连接),'#' 表示每层的末尾。

示例 2:

输入:root = []
输出:[]

提示:

  • 树中的节点数在范围 [0, 6000]
  • -100 <= Node.val <= 100

进阶:

  • 你只能使用常量级额外空间。
  • 使用递归解题也符合要求,本题中递归程序的隐式栈空间不计入额外空间复杂度。

Submission

运行时间: 64 ms

内存: 16 MB

"""
# Definition for a Node.
class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next
"""

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        if root is None:
            return
        if root.left is None and root.right is None:
            return root

        def find_next_right(root):
            while root.next:
                if root.next.left:
                    return root.next.left
                if root.next.right:
                    return root.next.right
                root = root.next

        if root.left:
            if root.right:
                root.left.next = root.right
                root.right.next = find_next_right(root)
            else:
                root.left.next = find_next_right(root)
        else:
            if root.right:
                root.right.next = find_next_right(root)
        self.connect(root.right)
        self.connect(root.left)
        return root

Explain

这个题解使用递归的思路来填充每个节点的 next 指针。主要分为以下几步: 1. 判断根节点是否为空或者是否为叶子节点,如果是则直接返回。 2. 定义一个辅助函数 find_next_right,用于查找当前节点的 next 指针应该指向的节点。该函数从当前节点的 next 节点开始,依次查找其左右子节点,直到找到一个非空的子节点或者 next 指针为空。 3. 如果当前节点有左子节点,则将其 next 指针指向右子节点(如果有)或者 find_next_right 函数的返回结果。 4. 如果当前节点没有左子节点但有右子节点,则将右子节点的 next 指针指向 find_next_right 函数的返回结果。 5. 递归处理当前节点的右子树和左子树。 6. 返回根节点。

时间复杂度: O(n^2)

空间复杂度: O(1)

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        if root is None:
            return
        if root.left is None and root.right is None:
            return root
        
        def find_next_right(root):
            while root.next:
                if root.next.left:
                    return root.next.left
                if root.next.right:
                    return root.next.right
                root = root.next
        
        if root.left:
            if root.right:
                root.left.next = root.right  # 将左子节点的 next 指针指向右子节点
                root.right.next = find_next_right(root)  # 将右子节点的 next 指针指向下一个右侧节点
            else:
                root.left.next = find_next_right(root)  # 如果没有右子节点,将左子节点的 next 指针指向下一个右侧节点
        else:
            if root.right:
                root.right.next = find_next_right(root)  # 如果没有左子节点,将右子节点的 next 指针指向下一个右侧节点
        
        self.connect(root.right)  # 递归处理右子树
        self.connect(root.left)  # 递归处理左子树
        
        return root

Explore

在递归处理时,先处理右子树再处理左子树确保了当处理左子节点时,它所有可能的右侧节点(属于当前节点的右子树)已经被正确连接。这是因为在填充每个节点的 next 指针时,左子节点可能需要连接到右子树的某个节点上。如果先递归处理左子树,那么右子树可能还未被处理,导致无法正确设置左子树中某些节点的 next 指针。

在 find_next_right 函数中,需要循环检查当前节点的 next 节点的左右子节点,因为二叉树可能不是完全二叉树,所以即使当前节点有 next 指针,其直接兄弟节点可能没有子节点。因此,需要遍历当前节点的所有后续兄弟节点,直到找到一个具有子节点的节点(无论是左子还是右子)。这确保了正确地找到下一个可用的子节点来连接。

如果 find_next_right 函数在遍历中遇到连续的 next 指针指向 NULL 的情况,函数将继续跟踪 next 指针直到达到链表的末尾(即当 next 指针为 NULL 时停止)。在这个过程中,如果找到了非空的子节点,则返回该节点;如果链表末尾仍未找到非空子节点,则返回 NULL。这保证了即使有多个连续的空 next 指针,也能正确处理或者确认没有可用的连接目标。

在递归解法中,如果遇到非均匀的二叉树,比如所有左子节点都不存在,算法仍然能有效地工作。每个节点只被访问一次以设置其 next 指针,因此时间复杂度仍然是 O(n),其中 n 是树中的节点数。虽然在极端情况下,比如所有节点都只有一个子节点,可能导致递归调用深度较深,但总体上算法的时间复杂性不会因树的形状而变差。不过,深度递归可能影响空间复杂度,尤其是在递归深度非常大时可能需要更多的栈空间。