2867. 统计树中的合法路径数目
题目描述：
给你一棵 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形成一棵合法的树。
思路：
1)虽然他是想构成一棵树，但个人觉得更像缩减版的图，不过最后一句提示“输入保证 edges 形成一棵合法的树。”保证了树的存在。但其实质相同，就是对图的每一个节点，深度遍历，然后统计“合法”的路径数量
2）枚举每个质数节点，从质数的邻居开始dfs,统计在不经过质数的前提下能访问到多少个非质数。以下图为例，假设2的邻居能访问到3，4，5个非质数。
4和左边这3个点，两两之间的路径都只包含质数2。5和左边这3+4个点，两两之间的路径都只包含质数2.根据乘法原理，把4*3+5*7加到答案中。注：只考虑左边是避免重复统计。最后，从2出发到下面这3+4+5=12个点的路径也只包含质数2把12加到答案中。
代码：
#标记10^5以内质数
MX=10**5+1
isPrime=[True]*MX
isPrime[1]=False
for i in range(2,isqrt(MX)+1):
if isPrime[i]:
for j in range(i*i,MX,i):#j为i的倍数,代表非质数
isPrime[j]=False
class Solution:
def countPaths(self, n: int, edges: List[List[int]]) -> int:
#构建邻接表
g=[[] for _ in range(n+1)]
for x,y in edges:
g[x].append(y)
g[y].append(x)
def dfs(x:int,fa:int)->None:
nodes.append(x)
for y in g[x]:
if y !=fa and not isPrime[y]:
dfs(y,x)
ans = 0
size = [0] * (n + 1)
for x in range(1, n + 1):
if not isPrime[x]: # 跳过非质数
continue
s=0
for y in g[x]: # 质数 x 把这棵树分成了若干个连通块
if isPrime[y]:
continue
if size[y]==0:#还没计算的
nodes=[]
dfs(y,-1) # 遍历 y 所在连通块，在不经过质数的前提下，统计有多少个非质数
for z in nodes:
size[z]=len(nodes)
# 这 size[y] 个非质数与之前遍历到的 s 个非质数，两两之间的路径只包含质数 x
ans+=size[y]*s
s+=size[y]
ans+=s #从x出发的路径
return ans
