LeetCode 题解工作台
统计树中的合法路径数目
给你一棵 n 个节点的无向树，节点编号为 1 到 n 。给你一个整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [u i , v i ] 表示节点 u i 和 v i 在树中有一条边。 请你返回树中的 合法路径数目 。 如果在节点 a 到节点 b 之间 …
5
题型
5
代码语言
3
相关题
当前训练重点
困难 · 二分·树·traversal
答案摘要
我们可以预处理得到 $[1, n]$ 中的所有质数，其中 表示 是否为质数。 接下来，我们根据二维整数整数构建图 ，其中 表示节点 的所有邻居节点。如果一条边的两个节点都不是质数，那么我们就将这两个节点合并到同一个连通分量中。
Interview AiBoxInterview AiBox 实时 AI 助手，陪你讲清 二分·树·traversal 题型思路
题目描述
给你一棵 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形成一棵合法的树。
解题思路
方法一：预处理 + 并查集 + 枚举
我们可以预处理得到 中的所有质数，其中 表示 是否为质数。
接下来，我们根据二维整数整数构建图 ，其中 表示节点 的所有邻居节点。如果一条边的两个节点都不是质数，那么我们就将这两个节点合并到同一个连通分量中。
然后，我们在 的范围内枚举所有质数 ，考虑包含 的所有路径。
由于 已经是质数，如果 是路径的一个端点，那么我们只需要累计与节点 相邻的所有连通分量的大小即可。如果 是路径上的某个中间点，那么我们需要累计相邻的任意两个连通分量的大小之积。
时间复杂度 ，空间复杂度 。其中 为节点数，而 为阿克曼函数的反函数。
class UnionFind:
def __init__(self, n):
self.p = list(range(n))
self.size = [1] * n
def find(self, x):
if self.p[x] != x:
self.p[x] = self.find(self.p[x])
return self.p[x]
def union(self, a, b):
pa, pb = self.find(a), self.find(b)
if pa == pb:
return False
if self.size[pa] > self.size[pb]:
self.p[pb] = pa
self.size[pa] += self.size[pb]
else:
self.p[pa] = pb
self.size[pb] += self.size[pa]
return True
mx = 10**5 + 10
prime = [True] * (mx + 1)
prime[0] = prime[1] = False
for i in range(2, mx + 1):
if prime[i]:
for j in range(i * i, mx + 1, i):
prime[j] = False
class Solution:
def countPaths(self, n: int, edges: List[List[int]]) -> int:
g = [[] for _ in range(n + 1)]
uf = UnionFind(n + 1)
for u, v in edges:
g[u].append(v)
g[v].append(u)
if prime[u] + prime[v] == 0:
uf.union(u, v)
ans = 0
for i in range(1, n + 1):
if prime[i]:
t = 0
for j in g[i]:
if not prime[j]:
cnt = uf.size[uf.find(j)]
ans += cnt
ans += t * cnt
t += cnt
return ans
复杂度分析
| 指标 | 值 |
|---|---|
| 时间 | and space complexity depend on the chosen approach. The sieve of Eratosthenes runs in O(n log log n), while the DFS traversal and dynamic programming state updates are O(n), making the overall time complexity approximately O(n log log n). Space complexity is O(n) due to the storage requirements for the sieve and dynamic programming table. |
| 空间 | Depends on the final approach |
面试官常问的追问
外企场景- question_mark
Check if the candidate can optimize tree traversal using DFS and dynamic programming.
- question_mark
Look for the ability to handle large inputs efficiently, using prime number computation and state tracking.
- question_mark
Evaluate how the candidate deals with the prime number constraint, especially in edge cases where there are few or no primes.
常见陷阱
外企场景- error
Forgetting to mark non-prime numbers correctly in the sieve of Eratosthenes, leading to incorrect path evaluations.
- error
Inefficient tree traversal that does not minimize redundant work, especially in larger trees.
- error
Incorrectly counting paths that contain more than one prime number, violating the problem's constraints.
进阶变体
外企场景- arrow_right_alt
Consider trees with only one node, where no paths exist.
- arrow_right_alt
Modify the problem to count paths with exactly two primes instead of one.
- arrow_right_alt
Explore how the solution scales when the tree is unbalanced or contains many leaves.