Count Valid Paths in a Tree
Count Valid Paths in a Tree
There is an undirected tree with n
nodes labeled from 1
to n
. You are given the integer n
and a 2D integer array edges
of length n - 1
, where edges[i] = [ui, vi]
indicates that there is an edge between nodes ui
and vi
in the tree.
Return the number of valid paths in the tree.
A path (a, b)
is valid if there exists exactly one prime number among the node labels in the path from a
to b
.
Note that:
- The path
(a, b)
is a sequence of distinct nodes starting with nodea
and ending with nodeb
such that every two adjacent nodes in the sequence share an edge in the tree. - Path
(a, b)
and path(b, a)
are considered the same and counted only once.
Example 1:
Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
Output: 4
Explanation: The pairs with exactly one prime number on the path between them are:
- (1, 2) since the path from 1 to 2 contains prime number 2.
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (2, 4) since the path from 2 to 4 contains prime number 2.
It can be shown that there are only 4 valid paths.
Example 2:
Input: n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
Output: 6
Explanation: The pairs with exactly one prime number on the path between them are:
- (1, 2) since the path from 1 to 2 contains prime number 2.
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (1, 6) since the path from 1 to 6 contains prime number 3.
- (2, 4) since the path from 2 to 4 contains prime number 2.
- (3, 6) since the path from 3 to 6 contains prime number 3.
It can be shown that there are only 6 valid paths.
Constraints:
1 <= n <= 105
edges.length == n - 1
edges[i].length == 2
1 <= ui, vi <= n
- The input is generated such that
edges
represent a valid tree.
解题思路
把简单路径按照经过的最高点进行分类,即根据路径两个端点的最近公共祖先来分类。对于以 为根的子树,假设与其相邻的子节点为 ,由于经过 的路径中必须恰好有一个质数节点,因此在子树 中,含端点 的路径中最多只能有一个质数节点。
对此定义 表示在子树 中含端点 的路径中恰好有 个质数节点的数量。当 是质数节点时,有 。当 不是质数节点时,有 。
再考虑两个端点的最近公共祖先是 且经过恰好一个质数节点的简单路径的数量。如果 是质数节点,则分成两个部分,一个是端点含 的,对应的数量是 。另一个两个端点来自不同的子树 ,对应的数量是 。如果 不是质数节点,也分成两个部分,一个是端点含 的,对应的数量是 。另一个两个端点来自不同的子树 ,对应的数量是 。
AC 代码如下,时间复杂度为 :
class Solution {
public:
long long countPaths(int n, vector<vector<int>>& edges) {
vector<vector<int>> g(n + 1);
for (auto &p : edges) {
g[p[0]].push_back(p[1]);
g[p[1]].push_back(p[0]);
}
vector<int> primes;
vector<bool> vis(n + 1);
vis[1] = true;
for (int i = 2; i <= n; i++) {
if (!vis[i]) primes.push_back(i);
for (int j = 0; primes[j] * i <= n; j++) {
vis[primes[j] * i] = true;
if (i % primes[j] == 0) break;
}
}
vector<vector<int>> f(n + 1, vector<int>(2));
long long ans = 0;
function<void(int, int)> dfs = [&](int u, int p) {
int s0 = 0, s1 = 0;
for (auto &v : g[u]) {
if (v == p) continue;
dfs(v, u);
s0 += f[v][0];
s1 += f[v][1];
}
if (!vis[u]) {
f[u][0] = 0;
f[u][1] = s0 + 1;
ans += s0;
for (auto &v : g[u]) {
if (v == p) continue;
s0 -= f[v][0];
ans += 1ll * f[v][0] * s0;
}
}
else {
f[u][0] = s0 + 1;
f[u][1] = s1;
ans += s1;
for (auto &v : g[u]) {
if (v == p) continue;
ans += 1ll * f[v][1] * (s0 - f[v][0]);
}
}
};
dfs(1, -1);
return ans;
}
};
本文来自博客园,作者:onlyblues,转载请注明原文链接:https://www.cnblogs.com/onlyblues/p/18064769
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 单线程的Redis速度为什么快?
· 展开说说关于C#中ORM框架的用法!
· Pantheons:用 TypeScript 打造主流大模型对话的一站式集成库
· SQL Server 2025 AI相关能力初探
· 为什么 退出登录 或 修改密码 无法使 token 失效