统计好节点的数目

2024-11-14

题目:

现有一棵 无向 树,树中包含 n 个节点,按从 0n - 1 标记。树的根节点是节点 0 。给你一个长度为 n - 1 的二维整数数组 edges,其中 edges[i] = [ai, bi] 表示树中节点 ai 与节点 bi 之间存在一条边。

如果一个节点的所有子节点为根的 子树包含的节点数相同,则认为该节点是一个 好节点

返回给定树中 好节点 的数量。子树 指的是一个节点以及它所有后代节点构成的一棵树。

示例 1:

输入:edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]

输出:7

说明:

img

树的所有节点都是好节点。

示例 2:

输入:edges = [[0,1],[1,2],[2,3],[3,4],[0,5],[1,6],[2,7],[3,8]]

输出:6

说明:

img

树中有 6 个好节点。上图中已将这些节点着色。

示例 3:

输入:edges = [[0,1],[1,2],[1,3],[1,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[9,12],[10,11]]

输出:12

解释:

img

除了节点 9 以外其他所有节点都是好节点。

提示:

  • 2 <= n <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • 输入确保 edges 总表示一棵有效的树。

思路:

首先根据边数组 edges 构建邻接表 g。在树中,边的数量为节点数量减 1。因此,n 为 edges 的长度加 1。再构造深度优先搜索,输入为当前遍历的节点 node 和其父节点 parent,返回值为以 node 为根节点的树的节点数量。需要递归调用 node 的所有子节点。因为 g 中存的是邻接关系,所以要跳过节点 parent。在计算节点数量和的同时,需要判断 node 的所有子节点是否拥有相同的节点数。如果是的话,将结果加 1。最后调用 dfs(0,−1) 并返回结果。

代码:

class Solution {
public:
    int countGoodNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        int res = 0;
        vector<vector<int>> graph(n);
        // 构建树, 由于我们可以通过从根节点0开始深搜,由于深搜天然携带父子关系,所以无向图也无所谓,
        for (const auto& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        // 传入 node 及其 父节点 paernt,返回以node为根的子树的节点数
        function<int(int, int)> dfs = [&](int node, int parent) -> int {
            bool valid = true;
            int treeSize = 0, subTreeSize = 0;
            for (int child : graph[node]) {
                if (child != parent) {
                    int size = dfs(child, node);
                    if (subTreeSize == 0) subTreeSize = size;
                    else if (subTreeSize != size) valid = false;  // 存在子节点为根的子树节点数不同
                    treeSize += size;
                }
            }
            if (valid) res++;
            return treeSize + 1;
        };
        dfs(0, -1);
        return res;
    }
};