之前是出去旅游了没发，现在开学了，继续每日一题，继续卷，一上来就是困难题😓，直接cv大法。
LeetCode
统计树中的合法路径数目
2867. 统计树中的合法路径数目 - 力扣（LeetCode）
题目描述
给你一棵 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)视为 同一条 路径，且只计入答案 一次 。
示例 1：
输入：n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
输出：4
解释：恰好有一个质数编号的节点路径有：
- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
只有 4 条合法路径。
示例 2：
输入：n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
输出：6
解释：恰好有一个质数编号的节点路径有：
- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
- (1, 6) 因为路径 1 到 6 只包含一个质数 3 。
- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
- (3, 6) 因为路径 3 到 6 只包含一个质数 3 。
只有 6 条合法路径。
提示：
1 <= n <= 105edges.length == n - 1edges[i].length == 21 <= ui, vi <= n- 输入保证
edges形成一棵合法的树。
思路
困难题，cv大法
代码
C++
const int MX = 1e5;
bool np[MX + 1]; // 质数=false 非质数=true
int init = []() {
np[1] = true;
for (int i = 2; i * i <= MX; i++) {
if (!np[i]) {
for (int j = i * i; j <= MX; j += i) {
np[j] = true;
}
}
}
return 0;
}();
class Solution {
public:
long long countPaths(int n, vector<vector<int>> &edges) {
vector<vector<int>> g(n + 1);
for (auto &e: edges) {
int x = e[0], y = e[1];
g[x].push_back(y);
g[y].push_back(x);
}
vector<int> size(n + 1);
vector<int> nodes;
function<void(int, int)> dfs = [&](int x, int fa) {
nodes.push_back(x);
for (int y: g[x]) {
if (y != fa && np[y]) {
dfs(y, x);
}
}
};
long long ans = 0;
for (int x = 1; x <= n; x++) {
if (np[x]) continue; // 跳过非质数
int sum = 0;
for (int y: g[x]) { // 质数 x 把这棵树分成了若干个连通块
if (!np[y]) continue;
if (size[y] == 0) { // 尚未计算过
nodes.clear();
dfs(y, -1); // 遍历 y 所在连通块，在不经过质数的前提下，统计有多少个非质数
for (int z: nodes) {
size[z] = nodes.size();
}
}
// 这 size[y] 个非质数与之前遍历到的 sum 个非质数，两两之间的路径只包含质数 x
ans += (long long) size[y] * sum;
sum += size[y];
}
ans += sum; // 从 x 出发的路径
}
return ans;
}
};
Java
class Solution {
private final static int MX = (int) 1e5;
private final static boolean[] np = new boolean[MX + 1]; // 质数=false 非质数=true
static {
np[1] = true;
for (int i = 2; i * i <= MX; i++) {
if (!np[i]) {
for (int j = i * i; j <= MX; j += i) {
np[j] = true;
}
}
}
}
public long countPaths(int n, int[][] edges) {
List<Integer>[] g = new ArrayList[n + 1];
Arrays.setAll(g, e -> new ArrayList<>());
for (var e : edges) {
int x = e[0], y = e[1];
g[x].add(y);
g[y].add(x);
}
long ans = 0;
int[] size = new int[n + 1];
var nodes = new ArrayList<Integer>();
for (int x = 1; x <= n; x++) {
if (np[x]) { // 跳过非质数
continue;
}
int sum = 0;
for (int y : g[x]) { // 质数 x 把这棵树分成了若干个连通块
if (!np[y]) {
continue;
}
if (size[y] == 0) { // 尚未计算过
nodes.clear();
dfs(y, -1, g, nodes); // 遍历 y 所在连通块，在不经过质数的前提下，统计有多少个非质数
for (int z : nodes) {
size[z] = nodes.size();
}
}
// 这 size[y] 个非质数与之前遍历到的 sum 个非质数，两两之间的路径只包含质数 x
ans += (long) size[y] * sum;
sum += size[y];
}
ans += sum; // 从 x 出发的路径
}
return ans;
}
private void dfs(int x, int fa, List<Integer>[] g, List<Integer> nodes) {
nodes.add(x);
for (int y : g[x]) {
if (y != fa && np[y]) {
dfs(y, x, g, nodes);
}
}
}
}
