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^5
• edges.length == n - 1
• edges[i].length == 2
• 1 <= 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