leetcode-第14周双周赛-1273-删除树节点

class Solution:
    def deleteTreeNodes(self, nodes: int, parent: List[int], value: List[int]) -> int:
        dp = [[0, 0]for _ in range(nodes)]
        p = collections.defaultdict(list)
        for i, v in enumerate(parent):
            p[v].append(i)
        for i in range(nodes - 1, -1, -1):
            dp[i][0] += value[i]
            dp[i][1] += 1
            for j in p[i]:
                dp[i][0] += dp[j][0]
                dp[i][1] += dp[j][1]
            if dp[i][0] == 0:
                dp[i][1] = 0
        return dp[0][1]

class Solution:
    def deleteTreeNodes(self, nodes: int, parent: List[int], value: List[int]) -> int:
        child = collections.defaultdict(list)
        for idx, par in enumerate(parent):
            child[par].append(idx)
        def get_count_val(i):
            count, val = 1, value[i]
            for c in child[i]:
                tc, tv = get_count_val(c)
                count += tc
                val += tv
            return (count, val) if val else (0, 0)
        return get_count_val(0)[0]