题目描述

给你一棵以节点 0 为根节点的树，定义如下：

• 节点的总数为 nodes 个
• 第 i 个节点的值为 value[i]
• 第 i 个节点的父节点是 parent[i]

请你删除节点值之和为 0 的每一棵子树。

在完成所有删除之后，返回树中剩余节点的数目。

样例

输入：nodes = 7, parent = [-1,0,0,1,2,2,2], value = [1,-2,4,0,-2,-1,-1]
输出：2

限制

• 1 <= nodes <= 10^4
• -10^5 <= value[i] <= 10^5
• parent.length == nodes
• parent[0] == -1 表示节点 0 是树的根。

算法

(深度优先遍历) $O(n)$

1. 通过 parent 数组建图，然后从根结点开始深度优先遍历。
2. 遍历的过程中返回当前子树的 sum 和结点的个数 tot。
3. 当前函数定义两个变量 sum 和 tot。初始时，sum = value[x]，tot = 1。
4. 如果子树返回的总和不为 0，则 sum 加上子树的 sum，tot 加上子树的 tot。
5. 递归结束前，如果 sum 为 0，则删除的结点累加上 tot。
6. 最后返回 nodes - tot。

时间复杂度

• 每个结点遍历一次，故时间复杂度为 $O(n)$。

空间复杂度

• 需要额外的空间建图，递归需要最多 $O(n)$ 的系统栈空间，故空间复杂度为 $O(n)$。

C++ 代码

class Solution {
public:
    vector<vector<int>> graph;
    int ans;

    pair<int, int> solve(int x, const vector<int> &value) {
        int sum = value[x];
        int tot = 1;

        for (int v : graph[x]) {
            auto t = solve(v, value);
            if (t.first != 0) {
                sum += t.first;
                tot += t.second;
            }
        }

        if (sum == 0)
            ans += tot;

        return make_pair(sum, tot);
    }

    int deleteTreeNodes(int nodes, vector<int>& parent, vector<int>& value) {
        graph.resize(nodes);

        for (int i = 1; i < nodes; i++)
            graph[parent[i]].push_back(i);

        ans = 0;
        solve(0, value);
        return nodes - ans;
    }
};