【Leetcode】1273. Delete Tree Nodes

题目地址：

https://leetcode.com/problems/delete-tree-nodes/

给定一棵树，顶点编号是$0\sim n-1$，$0$为树根，并且每个顶点对应一个权值，问删除所有权值和为$0$的子树后，整棵树还剩下多少节点。

思路是两遍DFS。第一遍DFS算出每个顶点为根的子树的权值和，以及每个顶点为根的子树的节点个数，第二遍DFS的时候看到权值和为$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)$。