从先序遍历还原二叉树

标签: 树 深度优先搜索 字符串 二叉树

难度: Hard

我们从二叉树的根节点 root 开始进行深度优先搜索。

在遍历中的每个节点处，我们输出 D 条短划线（其中 D 是该节点的深度），然后输出该节点的值。（如果节点的深度为 D，则其直接子节点的深度为 D + 1。根节点的深度为 0）。

如果节点只有一个子节点，那么保证该子节点为左子节点。

给出遍历输出 S，还原树并返回其根节点 root。

示例 1：

输入："1-2--3--4-5--6--7"
输出：[1,2,5,3,4,6,7]

示例 2：

输入："1-2--3---4-5--6---7"
输出：[1,2,5,3,null,6,null,4,null,7]

示例 3：

输入："1-401--349---90--88"
输出：[1,401,null,349,88,90]

提示：

  • 原始树中的节点数介于 1 和 1000 之间。
  • 每个节点的值介于 1 和 10 ^ 9 之间。

Submission

运行时间: 40 ms

内存: 16.4 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def recoverFromPreorder(self, traversal: str) -> Optional[TreeNode]:
        stack, idx = [], 0
        n = len(traversal)
        while idx < n:
            level, val = 0, ''
            # '-' 出现的次数代表当前节点的深度 level
            while idx < n and traversal[idx] == '-':
                level += 1
                idx += 1

            # 数字出现代表当前节点的值
            while idx < n and traversal[idx] != '-':
                val += traversal[idx]
                idx += 1
            
            # 确保当前节点所处的 level 与 节点深度对应
            while len(stack) > level:
                stack.pop()
            
            # 构建子节点，按照左右的顺序
            node = TreeNode(int(val))
            if stack and stack[-1].left is None:
                stack[-1].left = node
            elif stack and stack[-1].right is None:
                stack[-1].right = node
            # 不同 level 的放在不同的stack的索引
            stack.append(node)

        # print(len(stack), stack)
        return stack[0]

Explain

本题通过模拟树的深度优先搜索（DFS）的先序遍历来还原二叉树。首先解析输入的字符串，根据'-'的数量确定节点的深度，然后解析节点的数值。使用栈来存储已经解析的节点，并根据节点的深度调整栈的大小，以模拟回溯到正确的父节点。当解析到一个新节点时，根据其深度与栈顶节点的关系来确定它是左子节点还是右子节点。最后，栈底的节点是树的根节点。

时间复杂度: O(n)

空间复杂度: O(n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def recoverFromPreorder(self, traversal: str) -> Optional[TreeNode]:
        stack, idx = [], 0
        n = len(traversal)
        while idx < n:
            level, val = 0, ''
            # '-' 出现的次数代表当前节点的深度 level
            while idx < n and traversal[idx] == '-':
                level += 1
                idx += 1

            # 数字出现代表当前节点的值
            while idx < n and traversal[idx] != '-':
                val += traversal[idx]
                idx += 1
            
            # 确保当前节点所处的 level 与 节点深度对应
            while len(stack) > level:
                stack.pop()
            
            # 构建子节点，按照左右的顺序
            node = TreeNode(int(val))
            if stack and stack[-1].left is None:
                stack[-1].left = node
            elif stack and stack[-1].right is None:
                stack[-1].right = node
            # 不同 level 的放在不同的stack的索引
            stack.append(node)

        # 栈底元素即为根节点
        return stack[0]

Explore

在处理先序遍历字符串时，算法首先计数连续的'-'字符，这些'-'字符的数量表示当前节点的深度。完成对'-'的计数后，算法将继续读取直到遇到下一个'-'为止的字符序列，这部分字符序列代表当前节点的值。在这个过程中，使用两个循环确保深度和值的正确分离：第一个循环计算'-'的数量以确定深度，第二个循环读取数字来获取节点的值。

在构建新节点后将其添加到栈中的原因是栈在这里扮演了跟踪父节点和管理树结构的回溯过程的角色。栈顶的节点是当前正在处理的节点的父节点。当向树中添加新节点时，可以直接利用栈顶元素来确定新节点的父节点，并将新节点设置为父节点的左或右子节点。此外，栈的结构允许在遍历过程中容易地返回到之前的父节点，从而正确构造整个树的结构。

当栈的长度大于当前节点的深度时，弹出栈顶元素是为了确保回溯到正确的父节点。这与树的结构有着直接关系：在深度优先搜索中，当完成一个节点的所有子节点的添加后，需要返回到该节点的父节点以继续遍历或添加兄弟节点。栈的长度代表当前的节点深度，如果栈的长度大于新节点的深度，说明需要返回到更高的层级，因此弹出栈顶元素直到栈的长度等于新节点的深度，这样栈顶的节点就是新节点的父节点。

在代码中，如果一个新节点应该被添加为父节点的子节点时，首先检查父节点的左子节点是否为空。如果为空，则将新节点设置为左子节点。如果左子节点已存在，那么新节点将被添加到右子节点的位置（如果右子节点为空的话）。这种方式自然地处理了只有左子节点而没有右子节点的情况，因为新节点总是首先尝试被添加到左侧。只有当左子节点已经存在时，它才会被添加到右侧。