题目地址:
https://leetcode.com/problems/delete-tree-nodes/
给定一棵树,顶点编号是 0 ∼ n − 1 0\sim n-1 0∼n−1, 0 0 0为树根,并且每个顶点对应一个权值,问删除所有权值和为 0 0 0的子树后,整棵树还剩下多少节点。
思路是两遍DFS。第一遍DFS算出每个顶点为根的子树的权值和,以及每个顶点为根的子树的节点个数,第二遍DFS的时候看到权值和为 0 0 0的子树根的时候,就将总节点数减去那个子树的size。代码如下:
import java.util.*;
public class Solution {
public int deleteTreeNodes(int nodes, int[] parent, int[] value) {
// 邻接表建图
Map<Integer, List<Integer>> tree = buildTree(nodes, parent);
// numOfNodes[i]记录顶点编号为i的子树的节点总个数,初始化为1
int[] numOfNodes = new int[nodes];
Arrays.fill(numOfNodes, 1);
dfs(0, tree, value, numOfNodes);
return nodes - dfs2(0, tree, value, numOfNodes);
}
// 返回以cur为根的子树里的权值和为0的所有子树的节点个数总和
private int dfs2(int cur, Map<Integer, List<Integer>> tree, int[] value, int[] numOfNodes) {
// 如果cur子树权值就是0,则直接返回cur子树的节点个数
if (value[cur] == 0) {
return numOfNodes[cur];
}
// 否则递归累加其所有孩子里的权值和为0的子树总结点个数
int res = 0;
if (tree.containsKey(cur)) {
for (int next : tree.get(cur)) {
res += dfs2(next, tree, value, numOfNodes);
}
}
return res;
}
// 更新value和numOfNodes数组,使得value[i]是以i为树根的子树的权值和,numOfNodes[i]为以i为树根的子树的节点总个数
private void dfs(int cur, Map<Integer, List<Integer>> tree, int[] value, int[] numOfNodes) {
// 碰到叶子节点了,直接返回
if (!tree.containsKey(cur)) {
return;
}
for (int next : tree.get(cur)) {
// 更新next为根的子树的所有节点信息
dfs(next, tree, value, numOfNodes);
// 更新当前节点信息
value[cur] += value[next];
numOfNodes[cur] += numOfNodes[next];
}
}
private Map<Integer, List<Integer>> buildTree(int nodes, int[] parent) {
Map<Integer, List<Integer>> tree = new HashMap<>();
for (int i = 0; i < nodes; i++) {
tree.putIfAbsent(parent[i], new ArrayList<>());
tree.get(parent[i]).add(i);
}
return tree;
}
}
时空复杂度 O ( n ) O(n) O(n)。