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

题目来源

力扣每日一题；题序：2867

我的题解

方法一 埃氏筛+深度优先遍历

分别以质数节点为根，用「深度优先搜索」的方式，递归搜索所有的非质数的子树，并求出所有子树的大小，搜索过程中只搜索非质数节点。任何两个来自不同子树的节点，其路径都通过质数根节点，路径上恰好只有根节点一个质数节点，根据题意路径是合法的。只需要把所子树大小，两两相乘并求和，就可以得到包含根节点的所有合法路径。
具体步骤：

1. 使用欧拉筛计算每个节点是否是质数节点
2. 从质数节点开始进行深度优先遍历，直到遍历到另一个质数节点截止
3. 把所子树大小，两两相乘并求和

时间复杂度：O(n×loglogn)，其中 n是要筛查的数字范围。
空间复杂度：O(n)

class Solution {
    boolean[] isPrime;
    public long countPaths(int n, int[][] edges) {
        isPrime=new boolean[n+1];
        initPrime(n+1);
        List<Integer>[] g=createGraph(n,edges);
        long res=0;
        // 子树的大小
        long[] count=new long[n+1];

        List<Integer> visited=new ArrayList<>();
        for(int i=1;i<=n;i++){
            //需要从质数节点开始深搜
            if(!isPrime[i])
                continue;
            //记录与当前节点单独组成合法路径的节点数
            long cur=0;
            for(int j:g[i]){
                if(isPrime[j])
                    continue;
                if(count[j]==0) {
                    //每一个质数节点的已遍历节点需要重置
                    visited.clear();
                    //对于每个质数节点构造虚拟的前驱节点0
                    dfs(g, visited, j, 0);
                    //有多少个可以与当前质数节点构成合法路径的节点
                    long cnt = visited.size();
                    //
                    for (int k : visited)
                        count[k] = cnt;
                }
                // 当前子树与其他子树和的乘积
                res+=count[j]*cur;
                //更新所有子树和
                cur+=count[j];
            }
            res+=cur;
        }
        return res;
    }
    public void dfs(List<Integer>[] g,List<Integer> visited,int index,int pre){
        visited.add(index);
        for(int next:g[index]){
            //next!=pre防止循环遍历   !isPrime[next]邻接节点不是质数节点
            if(next!=pre&&!isPrime[next])
                dfs(g,visited,next,index);
        }
    }

    public List<Integer>[] createGraph(int n,int[][] edges){
        List<Integer>[] G=new ArrayList[n+1];
        for(int i=0;i<=n;i++){
            G[i]=new ArrayList<>();
        }
        for(int[] t:edges){
            int from=t[0];
            int to=t[1];
            G[from].add(to);
            G[to].add(from);
        }
        return G;
    }
    public void initPrime(int n){
        Arrays.fill(isPrime,true);
        isPrime[1]=false;
        for(int i=2;i*i<n;i++){
            if(isPrime[i]){
                for(int j=i*i;j<n;j+=i)
                    isPrime[j]=false;
            }
        }
    }

}

有任何问题，欢迎评论区交流，欢迎评论区提供其它解题思路（代码），也可以点个赞支持一下作者哈😄~