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

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

---

题目链接：2867. 统计树中的合法路径数目

代码如下：

//参考：https://leetcode.cn/problems/count-valid-paths-in-a-tree/solutions/2462627/shi-jian-ji-bai-100you-yu-suo-you-ti-jie-7e4s
const int N = 100001;
bool isPrime[N];

void initPrime() 
{
    fill(begin(isPrime),end(isPrime),true);
    if(!isPrime[10]) return;
    isPrime[0] = isPrime[1] = false;
    for(int i=2;i<N;i++) 
    {
        if(isPrime[i]) 
        {
            for(long long j=(long long)i*i;j<N;j+=i) 
                isPrime[j]= false;
        }
    }
}

class Solution {
public:
    void dfs(int& node,long long& count,vector<int>& pris,vector<vector<int>>& lnks,vector<bool>& visited)
    {
        visited[node]=true;
        for(auto&& i:lnks[node])
        {
            if(isPrime[i]) pris.emplace_back(i);
            else if(!visited[i])
                dfs(i,++count,pris,lnks,visited);
        }
    }

    long long countPaths(int n, const vector<vector<int>>& edges) 
    {
        initPrime();
        vector<vector<int>> lnks(n+1);
        for(const auto& edge:edges)
        {
            lnks[edge[0]].push_back(edge[1]);
            lnks[edge[1]].push_back(edge[0]);
        }
        vector<bool> visited(n+1,false);
        unordered_map<int,pair<long long,long long>> priCnt;
        vector<int> pris;
        long long count=0,num=0;
        // 深搜找出所有联通的非质数节点
        for(int i=1;i<=n;i++)
        {
            if(isPrime[i]||visited[i])// 质数 或 已访问 的节点直接跳过
                continue;
            num=1;
            pris.clear();
            dfs(i,num,pris,lnks,visited);
            count+=num*pris.size();
            for(auto &&pri:pris)
            {
                priCnt[pri].first+=num;
                priCnt[pri].second+=pow(num,2);
            }
        }

        for(auto&&[pri,nums]:priCnt)
        {
            count+=(pow(nums.first,2)-nums.second)/2;
        }
        return count;
    }
};