深搜+素数筛，LeetCode 2867. 统计树中的合法路径数目

一、题目

1、题目描述

给你一棵 n 个节点的无向树，节点编号为 1 到 n 。给你一个整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 在树中有一条边。

请你返回树中的 合法路径数目 。

如果在节点 a 到节点 b 之间 恰好有一个 节点的编号是质数，那么我们称路径 (a, b) 是 合法的 。

注意：

• 路径 (a, b) 指的是一条从节点 a 开始到节点 b 结束的一个节点序列，序列中的节点 互不相同 ，且相邻节点之间在树上有一条边。
• 路径 (a, b) 和路径 (b, a) 视为 同一条 路径，且只计入答案 一次 。

2、接口描述

​

class Solution {
public:
    long long countPaths(int n, vector<vector<int>>& edges) {
        
    }
};

3、原题链接

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

---

二、解题报告

1、思路分析

显然我们要能快速判断素数，所以在解题前先进行素数筛预处理

然后对于合法路径无非两种情况：
合数-素数-合数以及素数-合数

由于无向树中我们可以以任意节点为根，所以有上图

其中灰色三角形代表和质数节点直接相连的合数组成的连通块

那么含该素数的合法路径数目即为sz1 * sz2 + (sz1 + sz2) * sz3 + sz1 + sz2 + sz3

在无向树中我们可以通过深搜处理连通块情况，我们用sz[x]代表x所在合数连通块的节点数目

如果sz[x] = 0，说明该节点未处理，然后进行处理即可

2、复杂度

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

3、代码详解

​

const int N = 1e5 + 10;
bitset<N> isprime;//0质1合
int primes[10005], tot;
int init = [](){
    isprime[1] = 1;
    for(int i = 2; i <= 100000; i++){
        if(!isprime[i]) primes[tot++] = i;
        for(int j = 0; primes[j] <= 100000 / i; j++){
            isprime[primes[j] * i] = 1;
            if(i % primes[j] == 0)
                break;
        }    
    }
    return 0;
}();
struct edge{
    int v, nxt;
}edges[N << 1];
int head[N], idx;
void addedge(int u, int v){
    edges[idx] = {v , head[u]}, head[u] = idx ++;
}
class Solution {
public:
int sz[N];
    long long countPaths(int n, vector<vector<int>>& g) {
        memset(sz, 0, sizeof sz), memset(head, -1, sizeof head), idx = 0;
        for(auto& e : g) addedge(e[0], e[1]), addedge(e[1], e[0]);
        vector<int> cnt;
        function<void(int, int)> dfs = [&](int x, int fa){
            cnt.emplace_back(x);
            for(int i = head[x]; ~i; i = edges[i].nxt){
                int v = edges[i].v;
                if(v != fa && isprime[v])
                    dfs(v, x);
            }
        };
        long long ret = 0, s = 0;
        for(int i = 2; i <= n; i++){
            if(isprime[i]) continue;
            s = 0;
            for(int j = head[i]; ~j; j = edges[j].nxt){
                int v = edges[j].v;
                if(!isprime[v]) continue;
                if(!sz[v]) {
                    cnt.clear(), dfs(v, i);
                    for(auto x : cnt) sz[x] = cnt.size();
                }
                ret += s * sz[v], s += sz[v];
            }
            ret += s;
        }
        return ret;
    }
};