3249. 统计好节点的数目

来源第 410 场周赛 Q2难度中等分数1565

题目描述

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

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

返回给定树中 好节点 的数量。

子树指的是一个节点以及它所有后代节点构成的一棵树。

示例 1：

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

输出：7

说明：

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

示例 2：

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

输出：6

说明：

树中有 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

解释：

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

提示：

2 <= n <= 10⁵
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
输入确保 edges 总表示一棵有效的树。

解法

方法一：DFS

思考

好节点的所有子树大小相等。 \(n\le 10^5\)，对每个点重新算子树为平方。一次 DFS 即可在回溯时比较各子树大小并累加自身是否好。

\(\textit{dfs}(a,fa)\) 返回子树点数；若所有孩子返回值相同则 \(a\) 为好节点。无根树任取 \(0\) 为根。线性时间。

我们先根据题目给定的边 \(\textit{edges}\) 构建出树的邻接表 \(g\)，其中 \(g[a]\) 表示节点 \(a\) 的所有邻居节点。

然后，我们设计一个函数 \(\textit{dfs}(a, fa)\)，表示计算以节点 \(a\) 为根的子树中的节点数，并累计好节点的数量。其中 \(fa\) 表示节点 \(a\) 的父节点。

函数 \(\textit{dfs}(a, fa)\) 的执行过程如下：

初始化变量 \(\textit{pre} = -1\), \(\textit{cnt} = 1\), \(ok = 1\)，分别表示节点 \(a\) 的某个子树的节点数、节点 \(a\) 的所有子树的节点数、以及节点 \(a\) 是否为好节点。
遍历节点 \(a\) 的所有邻居节点 \(b\)，如果 \(b\) 不等于 \(fa\)，则递归调用 \(\textit{dfs}(b, a)\)，返回值为 \(\textit{cur}\)，并累加到 \(\textit{cnt}\) 中。如果 \(\textit{pre} < 0\)，则将 \(\textit{cur}\) 赋值给 \(\textit{pre}\)；否则，如果 \(\textit{pre}\) 不等于 \(\textit{cur}\)，说明节点 \(a\) 的不同子树的节点数不同，将 \(ok\) 置为 \(0\)。
最后，累加 \(ok\) 到答案中，并返回 \(\textit{cnt}\)。

在主函数中，我们调用 \(\textit{dfs}(0, -1)\)，最后返回答案。

时间复杂度 \(O(n)\)，空间复杂度 \(O(n)\)。其中 \(n\) 表示节点的数量。

Python3JavaC++GoTypeScript

class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        def dfs(a: int, fa: int) -> int:
            pre = -1
            cnt = ok = 1
            for b in g[a]:
                if b != fa:
                    cur = dfs(b, a)
                    cnt += cur
                    if pre < 0:
                        pre = cur
                    elif pre != cur:
                        ok = 0
            nonlocal ans
            ans += ok
            return cnt

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = 0
        dfs(0, -1)
        return ans

class Solution {
    private int ans;
    private List<Integer>[] g;

    public int countGoodNodes(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, k -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        dfs(0, -1);
        return ans;
    }

    private int dfs(int a, int fa) {
        int pre = -1, cnt = 1, ok = 1;
        for (int b : g[a]) {
            if (b != fa) {
                int cur = dfs(b, a);
                cnt += cur;
                if (pre < 0) {
                    pre = cur;
                } else if (pre != cur) {
                    ok = 0;
                }
            }
        }
        ans += ok;
        return cnt;
    }
}

class Solution {
public:
    int countGoodNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<int> g[n];
        for (const auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        int ans = 0;
        auto dfs = [&](this auto&& dfs, int a, int fa) -> int {
            int pre = -1, cnt = 1, ok = 1;
            for (int b : g[a]) {
                if (b != fa) {
                    int cur = dfs(b, a);
                    cnt += cur;
                    if (pre < 0) {
                        pre = cur;
                    } else if (pre != cur) {
                        ok = 0;
                    }
                }
            }
            ans += ok;
            return cnt;
        };
        dfs(0, -1);
        return ans;
    }
};

func countGoodNodes(edges [][]int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    var dfs func(int, int) int
    dfs = func(a, fa int) int {
        pre, cnt, ok := -1, 1, 1
        for _, b := range g[a] {
            if b != fa {
                cur := dfs(b, a)
                cnt += cur
                if pre < 0 {
                    pre = cur
                } else if pre != cur {
                    ok = 0
                }
            }
        }
        ans += ok
        return cnt
    }
    dfs(0, -1)
    return
}

function countGoodNodes(edges: number[][]): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    let ans = 0;
    const dfs = (a: number, fa: number): number => {
        let [pre, cnt, ok] = [-1, 1, 1];
        for (const b of g[a]) {
            if (b !== fa) {
                const cur = dfs(b, a);
                cnt += cur;
                if (pre < 0) {
                    pre = cur;
                } else if (pre !== cur) {
                    ok = 0;
                }
            }
        }
        ans += ok;
        return cnt;
    };
    dfs(0, -1);
    return ans;
}

方法二：显式栈

思考

好节点要求每个孩子的子树大小都相等。 \(n\) 可以到 \(10^5\)，对每个点重新遍历子树是平方时间。一条链上第一次递归总是走向唯一的孩子，深度就是 \(n\)，会超出默认递归上限。

子树大小要等孩子都处理完才能比较。

因此用显式栈做后序。进入节点时先压入离开标记，再压入子节点。离开时读取各孩子的子树大小，若它们相同则当前节点是好节点，并把自身子树大小写回。无根树以 \(0\) 为根。