过了5/10
,哪个佬帮忙优化掉一层循环
def dfs(u, kid): # 返回以kid为根节点的树的深度,并保存每层的节点,及每个节点属于哪一层
deep = 1
for son in c[kid]:
d[u].add(son) # d[u]:第u层所有点的集合
deep = max(deep, dfs(u + 1, son) + 1)
return deep
n = int(input())
p = [0, 0] + list(map(int, input().split())) # 记录每个点的父亲,根节点父亲设为0
w = [0] + list(map(int, input().split()))
c = [set() for _ in range(n + 1)] # 记录每个点的儿子
for i in range(2, n + 1): c[p[i]].add(i)
d = [set() for _ in range(n + 1)] # 记录每层的节点
d[1].add(1)
f = [[0, 0] for _ in range(n + 1)]
m = dfs(2, 1)
for t in range(m, 0, -1): # 遍历每一层,从低往高算,直到树根
for u in d[t]: # 计算每层所有的节点
for son in d[t + 1]:
# 不选u点
f[u][0] = max(max(f[son][0], f[son][1]), f[u][0])
# 选u点
f[u][1] = max(f[u][1], f[son][0])
if p[son] != u: # 如果son不是u的亲儿子,son可选
f[u][1] = max(f[son][1], f[u][1])
f[u][1] += w[u] # 最后加上u点权值
print(max(f[1][0], f[1][1]))