力扣每日一题统计树中的合法路径数目

最新推荐文章于 2024-10-01 05:04:32 发布

兑生

最新推荐文章于 2024-10-01 05:04:32 发布

阅读量500

点赞数 5

分类专栏： # 各平台OJ题解文章标签： leetcode 算法职场和发展

本文链接：https://blog.csdn.net/lt6666678/article/details/136320022

版权

各平台OJ题解专栏收录该内容

88 篇文章 0 订阅

订阅专栏

文章讲述了如何使用DFS算法解决2867问题，通过记忆化搜索避免重复计算非质数连通块，优化了时间复杂度，以求统计树中合法的质数路径数目。

摘要由CSDN通过智能技术生成

Problem: 2867. 统计树中的合法路径数目

思路

👨‍🏫 灵神题解
在这里插入图片描述
代码实现时，为避免反复 DFS 同一个非质数连通块，可以把每个非质数所处的连通块的大小记录下来（类似记忆化搜索）。如果之前计算过，就无需再次 DFS。

复杂度

时间复杂度: $O (n)$

空间复杂度: $O (n)$

Code

class Solution {
    private final static int MX = (int) 1e5;
    private final static boolean[] np = new boolean[MX + 1]; // 质数=false 非质数=true

    static {
        np[1] = true;
        for (int i = 2; i * i <= MX; i++) {
            if (!np[i]) {
                for (int j = i * i; j <= MX; j += i) {
                    np[j] = true;
                }
            }
        }
    }

    public long countPaths(int n, int[][] edges) {
        List<Integer>[] g = new ArrayList[n + 1];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x);
        }

        long ans = 0;
        int[] size = new int[n + 1];
        var nodes = new ArrayList<Integer>();
        for (int x = 1; x <= n; x++) {
            if (np[x]) { // 跳过非质数
                continue;
            }
            int sum = 0;
            for (int y : g[x]) { // 质数 x 把这棵树分成了若干个连通块
                if (!np[y]) {
                    continue;
                }
                if (size[y] == 0) { // 尚未计算过
                    nodes.clear();
                    dfs(y, -1, g, nodes); // 遍历 y 所在连通块，在不经过质数的前提下，统计有多少个非质数
                    for (int z : nodes) {
                        size[z] = nodes.size();
                    }
                }
                // 这 size[y] 个非质数与之前遍历到的 sum 个非质数，两两之间的路径只包含质数 x
                ans += (long) size[y] * sum;
                sum += size[y];
            }
            ans += sum; // 从 x 出发的路径
        }
        return ans;
    }

    private void dfs(int x, int fa, List<Integer>[] g, List<Integer> nodes) {
        nodes.add(x);
        for (int y : g[x]) {
            if (y != fa && np[y]) {
                dfs(y, x, g, nodes);
            }
        }
    }
}