Hard
题目描述
给你一棵 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 <= 10^5edges.length == n - 1edges[i].length == 21 <= ui, vi <= n- 输入保证
edges表示一棵有效的树。
解题思路
这是一道结合数论和树形动态规划的困难题目。核心思路是：对于每个节点作为最近公共祖先(LCA)，统计通过该节点且恰好包含一个质数的路径数。
解题步骤：
预处理质数：使用埃拉托色尼筛法找出1到n范围内的所有质数。
树形DP状态设计：对于每个节点i，定义：
dp[i][0]：从节点i开始向下的路径中，不含质数的路径数dp[i][1]：从节点i开始向下的路径中，恰好含1个质数的路径数
状态转移：
- 若节点i不是质数：
dp[i][0] = sum(dp[child][0]) + 1（自身也算一条路径）dp[i][1] = sum(dp[child][1])
- 若节点i是质数：
dp[i][0] = 0（质数节点开始的路径至少包含1个质数）dp[i][1] = sum(dp[child][0]) + 1
- 若节点i不是质数：
计算答案：对于每个节点作为LCA，统计通过它的有效路径：
- 如果当前节点是质数：计算来自不同子树的无质数路径的组合
- 如果当前节点不是质数：计算一个子树有1个质数、另一个子树无质数的路径组合
这种方法避免了暴力枚举所有路径，时间复杂度为O(n)。
代码实现
class Solution {
public:
long long countPaths(int n, vector<vector<int>>& edges) {
// 预处理质数
vector<bool> isPrime(n + 1, true);
isPrime[0] = isPrime[1] = false;
for (int i = 2; i * i <= n; i++) {
if (isPrime[i]) {
for (int j = i * i; j <= n; j += i) {
isPrime[j] = false;
}
}
}
// 构建邻接表
vector<vector<int>> adj(n + 1);
for (auto& edge : edges) {
adj[edge[0]].push_back(edge[1]);
adj[edge[1]].push_back(edge[0]);
}
// dp[i][0]: 从i开始不含质数的路径数
// dp[i][1]: 从i开始含1个质数的路径数
vector<vector<long long>> dp(n + 1, vector<long long>(2, 0));
long long result = 0;
function<void(int, int)> dfs = [&](int node, int parent) {
if (isPrime[node]) {
dp[node][0] = 0;
dp[node][1] = 1;
} else {
dp[node][0] = 1;
dp[node][1] = 0;
}
for (int child : adj[node]) {
if (child != parent) {
dfs(child, node);
if (isPrime[node]) {
// 当前节点是质数
result += dp[node][0] * dp[child][0];
dp[node][0] += dp[child][0];
} else {
// 当前节点不是质数
result += dp[node][0] * dp[child][1] + dp[node][1] * dp[child][0];
dp[node][0] += dp[child][0];
dp[node][1] += dp[child][1];
}
}
}
};
dfs(1, -1);
return result;
}
};
class Solution:
def countPaths(self, n: int, edges: List[List[int]]) -> int:
# 预处理质数
is_prime = [True] * (n + 1)
is_prime[0] = is_prime[1] = False
for i in range(2, int(n**0.5) + 1):
if is_prime[i]:
for j in range(i * i, n + 1, i):
is_prime[j] = False
# 构建邻接表
adj = [[] for _ in range(n + 1)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
# dp[i][0]: 从i开始不含质数的路径数
# dp[i][1]: 从i开始含1个质数的路径数
dp = [[0, 0] for _ in range(n + 1)]
result = 0
def dfs(node, parent):
nonlocal result
if is_prime[node]:
dp[node][0] = 0
dp[node][1] = 1
else:
dp[node][0] = 1
dp[node][1] = 0
for child in adj[node]:
if child != parent:
dfs(child, node)
if is_prime[node]:
# 当前节点是质数
result += dp[node][0] * dp[child][0]
dp[node][0] += dp[child][0]
else:
# 当前节点不是质数
result += dp[node][0] * dp[child][1] + dp[node][1] * dp[child][0]
dp[node][0] += dp[child][0]
dp[node][1] += dp[child][1]
dfs(1, -1)
return result
public class Solution {
public long CountPaths(int n, int[][] edges) {
// 预处理质数
bool[] isPrime = new bool[n + 1];
Array.Fill(isPrime, true);
isPrime[0] = isPrime[1] = false;
for (int i = 2; i * i <= n; i++) {
if (isPrime[i]) {
for (int j = i * i; j <= n; j += i) {
isPrime[j] = false;
}
}
}
// 构建邻接表
List<int>[] adj = new List<int>[n + 1];
for (int i = 0; i <= n; i++) {
adj[i] = new List<int>();
}
foreach (var edge in edges) {
adj[edge[0]].Add(edge[1]);
adj[edge[1]].Add(edge[0]);
}
// dp[i][0]: 从i开始不含质数的路径数
// dp[i][1]: 从i开始含1个质数的路径数
long[,] dp = new long[n + 1, 2];
long result = 0;
void Dfs(int node, int parent) {
if (isPrime[node]) {
dp[node, 0] = 0;
dp[node, 1] = 1;
} else {
dp[node, 0] = 1;
dp[node, 1] = 0;
}
foreach (int child in adj[node]) {
if (child != parent) {
Dfs(child, node);
if (isPrime[node]) {
// 当前节点是质数
result += dp[node, 0] * dp[child, 0];
dp[node, 0] += dp[child, 0];
} else {
// 当前节点不是质数
result += dp[node, 0] * dp[child, 1] + dp[node, 1] * dp[child, 0];
dp[node, 0] += dp[child, 0];
dp[node, 1] += dp[child, 1];
}
}
}
}
Dfs(1, -1);
return result;
}
}
var countPaths = function(n, edges) {
// 预处理质数
const isPrime = new Array(n + 1).fill(true);
isPrime[0] = isPrime[1] = false;
for (let i = 2; i * i <= n; i++) {
if (isPrime[i]) {
for (let j = i * i; j <= n; j += i) {
isPrime[j] = false;
}
}
}
// 构建邻接表
const adj = Array.from({length: n + 1}, () => []);
for (const [u, v] of edges) {
adj[u].push(v);
adj[v].push(u);
}
// dp[i][0]: 从i开始不含质数的路径数
// dp[i][1]: 从i开始含1个质数的路径数
const dp = Array.from({length: n + 1}, () => [0, 0]);
let result = 0;
function dfs(node, parent) {
if (isPrime[node]) {
dp[node][0] = 0;
dp[node][1] = 1;
} else {
dp[node][0] = 1;
dp[node][1] = 0;
}
for (const child of adj[node]) {
if (child !== parent) {
dfs(child, node);
if (isPrime[node]) {
// 当前节点是质数
result += dp[node][0] * dp[child][0];
dp[node][0] += dp[child][0];
} else {
// 当前节点不是质数
result += dp[node][0] * dp[child][1] + dp[node][1] * dp[child][0];
dp[node][0] += dp[child][0];
dp[node][1] += dp[child][1];
}
}
}
}
dfs(1, -1);
return result;
};
复杂度分析
| 复杂度类型 | 复杂度 | 说明 |
|---|---|---|
| 时间复杂度 | O(n log log n) | 埃拉托色尼筛法预处理质数 + 树的DFS遍历 |
| 空间复杂度 | O(n) | 存储邻接表、DP数组和质数标记数组 |