Algorithm/ct/l2/打家劫舍3.py

45 lines
1.5 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import sys
sys.setrecursionlimit(100010) # 修改默认递归深度
class BiTNode(object): # 二叉树节点
def __init__(self, data):
self.data = data
self.dp = [0, 0] # dp[0]表示不偷该节点,dp[1]表示偷该节点
self.lChild = None
self.rChild = None
# 二叉树的后序遍历
def dfs(T):
if T != None:
dfs(T.lChild)
dfs(T.rChild)
lChildDp = [0, 0]
rChildDp = [0, 0]
if T.lChild != None:
lChildDp = T.lChild.dp
if T.rChild != None:
rChildDp = T.rChild.dp
# 不偷该节点,左节点的最优情况(偷与不偷取大者) + 右节点的最优情况(偷与不偷取大者)
T.dp[0] = max(lChildDp[0], lChildDp[1]) + max(rChildDp[0], rChildDp[1])
# 偷该节点,该节点的现金 + 不偷左节点获得的现金 + 不偷右节点获得的现金
T.dp[1] = T.data + lChildDp[0] + rChildDp[0]
n = int(input())
nodes = list(map(int, input().split()))
fathers = list(map(int, input().split()))
nodeArray = [BiTNode(None) for i in range(n)]
for i in range(n): # 构造二叉树
nodeArray[i].data = nodes[i]
f = fathers[i] - 1 # 节点编号从0开始
if f >= 0:
if nodeArray[f].lChild == None:
nodeArray[f].lChild = nodeArray[i]
else:
nodeArray[f].rChild = nodeArray[i]
dfs(nodeArray[0]) # 二叉树的后序遍历nodeArray[0]为根节点
print(max(nodeArray[0].dp[0], nodeArray[0].dp[1]))