OD C卷 - 二叉树计算

二叉树计算

每个节点的值等于左子树、右子树中所有子孙节点的数值求和，构造出一个新的求和节点，整体结构与原始二叉树一致，只是节点data有变化。

输入描述
第一行输入二叉树的中序遍历；
第二行输入二叉树的前序遍历
输出描述
求和树的中序遍历，以空格分隔

示例1
输入：
7 -2 6 6 9
6 7 -2 9 6
输出：
-2 0 20 0 61

思路：

1. 根据输入恢复二叉树【递归】；
2. 构造求和树，并中序遍历【递归】；

# 节点定义
class Node:def __init__(self, data, left=None, right=None):self.data = dataself.left = leftself.right = rightdef create_btree(arr):""" 根据先序、中序遍历唯一恢复出二叉树 """if not arr:return None# 从先序遍历中取一个根global pre_travelif pre_travel:data = pre_travel.pop(0)root = Node(data)else:return None# 递归创建左右子树idx = arr.index(data)left_arr = arr[:idx]right_arr = arr[idx+1:]root.left = create_btree(left_arr)root.right = create_btree(right_arr)return rootmid_travel = list(map(int, input().strip().split()))
pre_travel = list(map(int, input().strip().split()))# 从前序遍历中依次取一个根节点，在中序遍历中一分为二，然后分别构建左右子树
root = create_btree(mid_travel)# 构建求和树
def calc_sum(root):if root is None:return 0elif root.left is None and root.right is None:return root.dataelse:return root.data + calc_sum(root.left) + calc_sum(root.right)def create_sum_tree(root):if root is None:return Noneroot.data = calc_sum(root.left) + calc_sum(root.right)root.left = create_sum_tree(root.left)root.right = create_sum_tree(root.right)return rootsum_root = create_sum_tree(root)# 中序遍历
result = []
def mid_travel(root):if root is None:returnmid_travel(root.left)result.append(root.data)mid_travel(root.right)returnmid_travel(sum_root)# 拼接输出字符串
output_str = ""
for e in result:output_str += str(e) + " "
print(output_str[:-1])