LC 2049. 统计最高分的节点数目

题目描述

这是 LeetCode 上的 2049. 统计最高分的节点数目 ，难度为 中等。

给你一棵根节点为 $0$ 的 二叉树 ，它总共有 $n$ 个节点，节点编号为 $0$ 到 $n - 1$ 。

同时给你一个下标从 $0$ 开始的整数数组 $parents$ 表示这棵树，其中 $parents[i]$ 是节点 $i$ 的父节点。由于节点 $0$ 是根，所以 $parents[0] == -1$ 。

一个子树的 大小 为这个子树内节点的数目。每个节点都有一个与之关联的 分数 。求出某个节点分数的方法是，将这个节点和与它相连的边全部 删除 ，剩余部分是若干个 非空 子树，这个节点的 分数 为所有这些子树 大小的乘积 。

请你返回有 最高得分 节点的 数目 。

示例 1:

输入：parents = [-1,2,0,2,0]

输出：3

解释：
- 节点 0 的分数为：3 * 1 = 3
- 节点 1 的分数为：4 = 4
- 节点 2 的分数为：1 * 1 * 2 = 2
- 节点 3 的分数为：4 = 4
- 节点 4 的分数为：4 = 4
最高得分为 4 ，有三个节点得分为 4 （分别是节点 1，3 和 4 ）。

示例 2：

输入：parents = [-1,2,0]

输出：2

解释：
- 节点 0 的分数为：2 = 2
- 节点 1 的分数为：2 = 2
- 节点 2 的分数为：1 * 1 = 1
最高分数为 2 ，有两个节点分数为 2 （分别为节点 0 和 1 ）。

提示：

• $n == parents.length$
• $2 <= n <= 10^5$
• $parents[0] == -1$
• 对于 $i != 0$ ，有 $0 <= parents[i] <= n - 1$
• $parents$ 表示一棵二叉树。

---

建图 + DFS + 统计答案

为了更具有一般性，我们假定该树为多叉树。

由于题目给定的 parents 数组仅支持我们快速查找某个节点的父节点，为了方便遍历整棵树，我们先使用「邻接表」将图（树）建起来。

还不熟悉「邻接表」存图方式的同学可以看看前置 🧀 : 涵盖所有的「存图方式」模板。

然后使用 DFS 预处理出 f 数组，其中 $f[i]$ 代表以节点 $i$ 为根节点的子树所包含的节点数量。

考虑如何计算「删除某个节点 $x$ 后，剩余连通块的数量，以及每个连通块的节点数量」，根据节点 $x$ 是否为根节点进行分情况讨论：

• 若 $x$ 为根节点，删除后的连通块的数量为「$x$ 的出边数量」，假定共有 $k$ 条出边，根据题目定义，对应的 大小 为各个连通块的节点数量乘积：

  $$f[u_1] \times f[u_2] \times ... \times f[u_k]$$

• 若 $x$ 不是根节点，删除后的连通块的数量为「$x$ 的出边数量 + $1$」，其中 $1$ 代指「以 $x$ 节点的父节点所在的整体连通块」。

  假定节点 $x$ 共有 $k$ 条出边，根据题目定义，对应的 大小 为「(各个连通块的节点数量乘积) $\times$ ($x$ 节点的父节点所在的整体连通块大小)」，而「$x$ 节点的父节点所在的整体连通块大小」，利用容斥原理可知为 $f[root] - f[u] = n - f[u]$，含义为「从原树中减掉以节点 $x$ 为根节点的子树」的部分，即最终 大小 为：

  $$(f[u_1] \times f[u_2] \times ... \times f[u_k]) \times (n - f[x])$$

代码：

class Solution {
    static int N = 100010, M = N * 2;
    static int[] he = new int[N], e = new int[M], ne = new int[M];
    static int[] f = new int[N];
    int idx;
    void add(int a, int b) {
        e[idx] = b;
        ne[idx] = he[a];
        he[a] = idx++;
    }
    public int countHighestScoreNodes(int[] parents) {
        Arrays.fill(he, -1);
        int n = parents.length;
        for (int i = 1; i < n; i++) add(parents[i], i);
        dfs(0);
        long max = 0;
        int ans = 0;
        for (int x = 0; x < n; x++) {
            long cur = 1;
            for (int i = he[x]; i != -1; i = ne[i]) cur *= f[e[i]];
            if (x != 0) cur *= n - f[x];
            if (cur > max) {
                max = cur; ans = 1;
            } else if (cur == max) {
                ans++;
            }
        }
        return ans;
    }
    int dfs(int u) {
        int ans = 1;
        for (int i = he[u]; i != -1; i = ne[i]) ans += dfs(e[i]);
        f[u] = ans;
        return ans;
    }
}

• 时间复杂度：建图复杂度为 $O(n)$；通过 DFS 预处理 f 数组复杂度为 $O(n + m)$，其中 $m$ 为边数，对于本题（二叉树）而言，点边数量级相等，因此 DFS 预处理的复杂度为 $O(n)$；模拟删除任意点并统计答案的复杂度为 $O(n + m)$，对于本题（二叉树）而言，数量级为 $O(n)$。整体复杂度为 $O(n)$
• 空间复杂度：$O(n)$

---

拓扑排序 + 统计答案

通过对「待删除节点是否为根节点」的分情况讨论可知：若要算得删除某个节点后的得分，重点需要知道「当前节点的左右子树（如果有）所包含的节点数量」以及「当前节点的父节点所在连通块的节点数量」。

我们解法 $1$ 的建图目的就是为了得到某个节点的子树情况，DFS 预处理 f 数组的目的就是知道以某个节点为根的子树所包含的总节点数量。

而这个「建图 + DFS」过程可由「拓扑排序」所代替。

具体的，我们可以先对 parents进行遍历， 统计所有节点的出度 $out[i]$，然后将所有出度为 $0$ 的节点（叶子节点）进行入队。跑一遍求拓扑排序的 BFS，每次某个节点 $t$ 出队，我们就对节点 $t$ 的父节点 $fa = parents[t]$ 进行出度减一操作（若出度减至 $0$，则将 $fa$ 进行入队，注意 $fa$ 不能为根节点，因为 $fa$ 入队没有意义，其入队不能更新其他点的出度），并在求拓扑序的过程中预处理出 l 和 r 数组，$l[i]$ 和 $r[i]$ 分别代表节点 $i$ 的左子树的节点数和右节点的节点数。

跑完拓扑排序后，我们得到的 l 和 r 数组就足够我们统计答案，仍然是对删除节点 $x$ 是否为根节点的分情况讨论：

• 若 $x$ 不是根节点，得分为 $\max(l[x], 1) \times \max( r[x], 1)$
• 若 $x$ 为根节点，得分为 $\max(l[x], 1) \times \max( r[x], 1) \times (n - (l[x] + r[x] + 1)$

代码：

class Solution {
    public int countHighestScoreNodes(int[] parents) {
        int n = parents.length;
        int[] out = new int[n];
        for (int i = 1; i < n; i++) out[parents[i]]++;
        Deque<Integer> d = new ArrayDeque<>();
        for (int i = 0; i < n; i++) {
            if (out[i] == 0) d.addLast(i);
        }
        // l[i] 和 r[i] 分别代表节点 i 的左子树的节点数和右节点的节点数
        int[] l = new int[n], r = new int[n];
        while (!d.isEmpty()) {
            int t = d.pollFirst(), fa = parents[t];
            out[fa]--;
            if (l[fa] == 0) l[fa] = l[t] + r[t] + 1;
            else r[fa] = l[t] + r[t] + 1;
            if (out[fa] == 0 && fa != 0) d.addLast(fa);
        }
        long max = 0;
        int ans = 0;
        for (int i = 0; i < n; i++) {
            long cur = Math.max(l[i], 1) * Math.max(r[i], 1);
            if (i != 0) cur *= n - (l[i] + r[i] + 1);
            if (cur > max) {
                max = cur; ans = 1;
            } else if (cur == max) {
                ans++;
            }
        }
        return ans;
    }
}

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.2049 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。