首页手记 LeetCode 1028. 从先序遍历还原二叉树 |...

LeetCode 1028. 从先序遍历还原二叉树 | Python

标签：

Python 算法

1028. 从先序遍历还原二叉树

题目来源：力扣（LeetCode）https://leetcode-cn.com/problems/recover-a-tree-from-preorder-traversal

题目

我们从二叉树的根节点 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 之间。

解题思路

思路：栈 + 迭代

根据题意，结合示例和图例，我们可以得到如下信息：

遍历字符串时，读取 - 字符，直至遇到非 - 字符。此时我们可以通过 - 的个数判断当前节点的深度。（也印证题目中所说【在遍历中的每个节点处，我们输出 D 条短划线（其中 D 是该节点的深度）】）；
读取数字（这里注意示例 3，数字不仅是个位数），直到遇到非数字。这些数字则是节点的值。

在题目中提及【如果节点只有一个子节点，那么保证该子节点为左子节点。】，根据这个前提开始进行分析。假设当前节点为 A，上一个节点为 B，那么这里可能出现的就是两种情况：

A 是 B 的左子节点；
A 不是 B 的左子节点；

第一种情况就不解释了，这个相当于考虑前面的前提，因为当节点只有一个子节点，这个节点优先考虑是左子节点。所以这里先构建左子树。

因为本篇幅使用栈来辅助解决问题，栈存储的是等待构建子树的节点，当子树构建完成时，出栈。

如果当前节点的深度 < 栈的 size 时，这就表明上一个节点并不是当前节点的左子节点，也就是第二种所述的情况，（根据前面题目所提及的，如果节点只有一个子节点，这个子节点是左子节点）那么此时的 A 节点就是根节点到 B 节点（除 B 节点）这条路径的某个节点的右子节点。（这里考虑栈的 size 和当前节点的深度。）

具体的实现代码如下。（含注释）

代码实现

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def recoverFromPreorder(self, S: str) -> TreeNode:
        # 使用栈辅助
        stack = []
        
        s_length = len(S)

        i = 0
        while i < s_length:
            # 初始化深度，每个节点对应单独 cur_level
            cur_level = 0

            # 遍历查找 '-' 符号，直至遇到非 '-' 符号，'-' 的个数表示当前节点的深度
            while i < s_length and S[i] == '-':
                i += 1
                cur_level += 1
            
            # 记录节点的值
            value = 0
            while i < s_length and S[i] != '-':
                # 这里是因为节点值不一定是个位数，例如示例 3
                value = value * 10 + (ord(S[i]) - ord('0'))
                i += 1

            # 构建树
            node = TreeNode(value)

            # 先考虑根节点入栈
            if not stack:
                stack.append(node)
                continue

            if len(stack) == cur_level:
                # 先构建左子树
                stack[-1].left = node
            else:
                # 考虑到可能出现文章中所述的第二种情况，这里先进行判断，
                # 当前节点深度小于栈的 size 时，表示栈顶节点不是当前节点的父节点，出栈
                while cur_level < len(stack):
                    stack.pop()
                if len(stack) == cur_level:
                    # 构建右子树，因为左子树已经优先构建了
                    stack[-1].right = node
            # 这里需要将节点入栈，后续构建子树
            stack.append(node)

        return stack[0]

实现结果

总结

根据题意，结合示例和图例，可以得到信息如下：
- 遍历字符串，查找 - 字符，直至遇到非 - 字符，这里 - 字符的个数表示当前节点的深度；
- 遍历字符串，查找数字，直至遇到非数字。这里数字就是当前节点的值
使用栈辅助，栈中存储的是待构建子树的节点。假设当前节点为 A，上个节点为 B，那么可能会出现的情况如下：
- A 是 B 的左子节点
- A 不是 B 的左子节点
  第一种情况，根据题意【如果节点只有一个子节点，那么保证该子节点为左子节点】，这里不展开解释。第二种情况，如果 A 不是 B 的左子节点，再加上【如果节点只有一个子节点，那么保证该子节点为左子节点】，那么 A 可能是根节点到 B 节点（除 B 节点外）某个节点的右子节点（这里在代码中由栈的 size 与当前节点的深度的比较中体现）