一、题目
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;}
};