LeetCode1028 从先序遍历还原二叉树

题目描述

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

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

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

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

题解

我的方法（使用栈）

思路

1. 使用 itertools 库中的 groupby 函数对输入字符串 S 进行划分，得到列表 S_list 。如输入 S=“1-2–3--4-5–6--7” 被划分为 S_list=[1, ‘-’, 2, ‘–’, 3, ‘–’, 4, ‘-’, 5, ‘–’, 6, ‘–’, 7] 。其中， ‘-’ 字符串元素（如 S_list[1]）的长度表示当前节点的深度，数值元素（如 S_list[2]）表示当前节点的值。
2. 使用栈 stack 来还原树，并使用字典 d 来记录所有结点的深度：
   • 先创建根节点 root，将其压栈，并记录其深度 d[root]=0。
   • 对列表中的每一组 字符串-数值元素对，创建当前节点 pre_node ，并计算当前节点的深度 deep 。比较栈顶元素的深度 d[stack[-1]] 与 当前节点的深度 deep ，如果 d[stack[-1]] >= deep ，则弹出栈顶元素。重复该操作直至 d[stack[-1]] < deep。此时，栈顶节点 stack[-1] 是当前节点 pre_node 的父节点，将当前节点 pre_node 插入树中。

复杂度

时间复杂度

$O(n)$
其中，$n$是树中节点的个数。

空间复杂度

$O(n)$
其中，$n$是树中节点的个数。

代码

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

class Solution(object):
    def recoverFromPreorder(self, S):
        """
        :type S: str
        :rtype: TreeNode
        """
        from itertools import groupby
        S_list = [int(''.join(list(g))) if k else ''.join(list(g)) for k, g in groupby(S, key=lambda x: x.isdigit())]
        # [1, '-', 2, '--', 3, '--', 4, '-', 5, '--', 6, '--', 7]

        def insert_node(parent, child):
            if parent.left == None:
                parent.left = child
                return True
            elif parent.right == None:
                parent.right = child
                return True
            else:
                return False

        stack = []
        d = {} # 记录节点的深度
        root = TreeNode(S_list[0])
        stack.append(root)
        d[root] = 0
        for i in range(1, len(S_list), 2): # S_list[i] = '-'
            pre_node = TreeNode(S_list[i+1])
            deep = len(S_list[i])
            d[pre_node] = deep
            while d[stack[-1]] >= deep:
                stack.pop()
            insert_node(stack[-1], pre_node)
            stack.append(pre_node)
        return root