【每日一题】建图 + DFS 运用题

宫水三叶的刷题日记

发布于 2022-04-07 06:45:17

63000

代码可运行

文章被收录于专栏：宫水三叶的刷题日记宫水三叶的刷题日记

运行总次数：0

代码可运行

题目描述

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

Tag : 「图论」、「线性 DP」、「容斥原理」

给你一棵根节点为 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)