【leetcode】1273. Delete Tree Nodes

题目如下：

A tree rooted at node 0 is given as follows:

• The number of nodes is nodes;
• The value of the i-th node is value[i];
• The parent of the i-th node is parent[i].

Remove every subtree whose sum of values of nodes is zero.

After doing so, return the number of nodes remaining in the tree.

Example 1:

Input: nodes = 7, parent = [-1,0,0,1,2,2,2], value = [1,-2,4,0,-2,-1,-1]
Output: 2

Constraints:

• 1 <= nodes <= 10^4
• -10^5 <= value[i] <= 10^5
• parent.length == nodes
• parent[0] == -1 which indicates that 0 is the root.

解题思路：我的方法是递归+动态规划。对于任意一个节点i，记dp[i]为其子树的和，如果j,k....n为其子节点，那么有dp[i] = dp[j] + dp[k] + .... + dp[n] + value[i]。通过递归的方式很容易可以求出每个节点的子树和，相应的可以求出哪些节点的子树和为0。再遍历这些子树的所有节点，并标记为删除的状态，最后统计出状态为删除的节点即可。

代码如下：

class Solution(object):
    def deleteTreeNodes(self, nodes, parent, value):
        """
        :type nodes: int
        :type parent: List[int]
        :type value: List[int]
        :rtype: int
        """
        import sys
        sys.setrecursionlimit(1000000)
        dic = {}
        for i in range(len(parent)):
            dic[parent[i]] = dic.setdefault(parent[i],[]) + [i]
        dp = [None] * nodes
        def getValue(inx):
            if inx not in dic:
                dp[inx] = value[inx]
                return value[inx]
            elif dp[inx] != None:
                return dp[inx]
            count = 0
            for child in dic[inx]:
                count += getValue(child)
            count += value[inx]
            dp[inx] = count
            return count

        dic_remove = {}

        for i in range(nodes):
            if dp[i] == None:
                dp[i] = getValue(i)
            if dp[i] == 0: dic_remove[i] = 1

        delete = [0] * nodes

        def markDelete(inx):
            delete[inx] = 1
            if inx not in dic:
                return
            for key in dic[inx]:
                markDelete(key)

        for inx in dic_remove.iterkeys():
            markDelete(inx)
        res = sum(delete)
        return nodes - res