比较好理解的并查集做法
题目描述
我们要求的实际上是 非质数-质数-非质数 这样的序列个数(如果有一个非质数的部分为空则为特殊情况,即非质数-质数)
首先我们先用质数筛筛出1e5之内的质数,之后我们开始建图并维护一个并查集,其中质数点的权值为0,非质数点的权值为1,当我们遇到一条连接两个非质数的边时则把对应的两个集合合并。最后我们遍历所有的质数作为中间点,然后求出以当前质数点为中心的方案个数,累加即可。注意这里,如果我们二重循环遍历枚举所有的和当前质数点有边连接的非质数的话,时间复杂度最差为$O(n^2)$,即菊花图,这会超时,所以我们先遍历一遍计算出和当前质数点相连的非质数点的集合的大小总和,然后在再遍历一遍, 这次我们在总和中刨除当前的大小后计算合法的组合数,这样会将时间复杂度降到$O(n)$。
(但是要注意我们每一对算了两遍所以要除以2)
除此之外就是上文提到的特殊情况, 即质数点作为边上一个点的情况, 我们直接加上一个刚刚计算的相连的非质数点集合的大小总和sum即可, 具体实现见代码
时间复杂度
假如我们视并查集的时间复杂度为常数,建图的时间复杂度为$O(n)$
我们枚举每一个质数, 同时枚举它的边, 如果是图的话, 时间复杂度为$O(n+m)$
其中n为点数, m为边数, 由于本题是树, $m = n - 1$
所以时间复杂度为 $O(n)$
并查集代码
const int N = 100010;
int primes[N], cnt = 0;
bool st[N];
struct DSU {
std::vector<int> f, siz;
DSU(int n) : f(n), siz(n) {
std::iota(f.begin(), f.end(), 0);
for (int i = 1; i < n; i ++) {
siz[i] = st[i];
}
}
int leader(int x) {
while (x != f[x]) x = f[x] = f[f[x]];
return x;
}
bool same(int x, int y) { return leader(x) == leader(y); }
bool merge(int x, int y) {
x = leader(x);
y = leader(y);
if (x == y) return false;
siz[x] += siz[y];
f[y] = x;
return true;
}
int size(int x) { return siz[leader(x)]; }
};
class Solution {
public:
using ll = long long;
void get_primes() {
if (cnt > 0) return;
const int n = 100000;
for (int i = 2; i <= n; i ++ ) {
if (!st[i]) primes[cnt ++ ] = i;
for (int j = 0; primes[j] <= n / i; j ++ ) {
st[primes[j] * i] = true;
if (i % primes[j] == 0) break;
}
}
st[0] = st[1] = true;
}
long long countPaths(int n, vector<vector<int>>& edges) {
get_primes();
// 由于LC独特的计时方式,我们预处理一下,所有的测试用例只跑一次质数筛
vector<vector<int>> g(n + 1);
DSU dsu(n + 1);
for (auto& e : edges) {
int a = e[0], b = e[1];
g[a].push_back(b);
g[b].push_back(a);
if (st[a] and st[b]) dsu.merge(a, b);
}
ll res = 0;
// 枚举质数作为中心点
for (int i = 2; i <= n; i ++) {
if (st[i]) continue;
ll sum = 0, ans = 0;
for (auto x : g[i]) {
sum += dsu.size(x);
}
for (auto x : g[i]) {
int t = dsu.size(x);
ans += (sum - t) * t;
}
res += sum + ans / 2;
}
return res;
}
};