You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
hello-algo/codes/python/chapter_divide_and_conquer/build_tree.py

55 lines
1.4 KiB

"""
File: build_tree.py
Created Time: 2023-07-15
Author: krahets (krahets@163.com)
"""
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))
from modules import TreeNode, print_tree
def dfs(
preorder: list[int],
inorder_map: dict[int, int],
i: int,
l: int,
r: int,
) -> TreeNode | None:
"""构建二叉树:分治"""
# 子树区间为空时终止
if r - l < 0:
return None
# 初始化根节点
root = TreeNode(preorder[i])
# 查询 m ,从而划分左右子树
m = inorder_map[preorder[i]]
# 子问题:构建左子树
root.left = dfs(preorder, inorder_map, i + 1, l, m - 1)
# 子问题:构建右子树
root.right = dfs(preorder, inorder_map, i + 1 + m - l, m + 1, r)
# 返回根节点
return root
def build_tree(preorder: list[int], inorder: list[int]) -> TreeNode | None:
"""构建二叉树"""
# 初始化哈希表,存储 inorder 元素到索引的映射
inorder_map = {val: i for i, val in enumerate(inorder)}
root = dfs(preorder, inorder_map, 0, 0, len(inorder) - 1)
return root
"""Driver Code"""
if __name__ == "__main__":
preorder = [3, 9, 2, 1, 7]
inorder = [9, 3, 1, 2, 7]
print(f"前序遍历 = {preorder}")
print(f"中序遍历 = {inorder}")
root = build_tree(preorder, inorder)
print("构建的二叉树为:")
print_tree(root)