We can use a post-order DFS to compute the sum of values and the size (number of nodes) for each subtree. If a subtree’s sum is zero, we remove it (do not count its nodes). Otherwise, we keep its nodes. This approach efficiently traverses the tree and handles deletions in a single pass.
classSolution {
publicintdeleteTreeNodes(int nodes, int[] parent, int[] value) {
List<Integer>[] g =new List[nodes];
for (int i = 0; i < nodes; i++) g[i]=new ArrayList<>();
for (int i = 1; i < nodes; i++) g[parent[i]].add(i);
int[] ans = dfs(0, g, value);
return ans[1];
}
privateint[]dfs(int u, List<Integer>[] g, int[] value) {
int sum = value[u], sz = 1;
for (int v : g[u]) {
int[] res = dfs(v, g, value);
sum += res[0];
sz += res[1];
}
if (sum == 0) returnnewint[]{0, 0};
returnnewint[]{sum, sz};
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
classSolution {
fundeleteTreeNodes(nodes: Int, parent: IntArray, value: IntArray): Int {
val g = Array(nodes) { mutableListOf<Int>() }
for (i in1 until nodes) g[parent[i]].add(i)
fundfs(u: Int): Pair<Int, Int> {
var sum = value[u]
var sz = 1for (v in g[u]) {
val(s, cnt) = dfs(v)
sum += s
sz += cnt
}
returnif (sum ==0) 0 to 0else sum to sz
}
return dfs(0).second
}
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
classSolution:
defdeleteTreeNodes(self, nodes: int, parent: list[int], value: list[int]) -> int:
g: list[list[int]] = [[] for _ in range(nodes)]
for i in range(1, nodes):
g[parent[i]].append(i)
defdfs(u: int) -> tuple[int, int]:
s, sz = value[u], 1for v in g[u]:
cs, csz = dfs(v)
s += cs
sz += csz
if s ==0:
return0, 0return s, sz
return dfs(0)[1]