LC 2867. 统计树中的合法路径数目

https://leetcode.cn/problems/count-valid-paths-in-a-tree/description/

这题感觉自己是魔怔了，回过头来看其实是挺简单的题目的。

我一度有过这样的想法，可能是之前有过类似的题目

计算每个节点x，从root到该节点上，有多少个质数，假设有C个
然后计算这个的同时，考虑[C-2,C+2]这些节点的个数，两者相减（其实这里错误的离谱，不能相减啊!!!）
相减完了之后需要补回两个节点的LCA. 这个过程正好复习一下tarjan lca的算法。

最后提交的程序完全不是那么回事。对于这类枚举量比较小的情况，应该是只需要考虑每个子树然后拼凑起来就行

假设每个子树路径上有c0个没有任何质数的path, 有c1个只有1个质数的path.
观察当前子树，假设是(a, b)
如果父节点是质数的话，那么两边子树只能取没有质数的path, 就是 c0 * a
如果父节点不是质数的话，那么就是 c0 * b + c1 * a
向上返回的时候，也需要考虑父节点是否是质数。

def get_primes(N):
    ps = []
    mask = [0] * (N + 1)
    for i in range(2, N + 1):
        if mask[i] == 1: continue
        for j in range(2, N + 1):
            if i * j > N: break
            mask[i * j] = 1
    for i in range(2, N + 1):
        if mask[i] == 0:
            ps.append(i)
    return ps


class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        primes = set(get_primes(n))
        adj = [[] for _ in range(n + 1)]
        for x, y in edges:
            adj[x].append(y)
            adj[y].append(x)

        ans = 0

        def dfs(x, p):
            nonlocal ans
            isPrime = x in primes

            c0, c1 = 0, 0
            for y in adj[x]:
                if y == p: continue
                a, b = dfs(y, x)
                if isPrime:
                    ans += c0 * a
                else:
                    ans += c0 * b + c1 * a
                c0 += a
                c1 += b

            if isPrime:
                ans += c0
                res = (0, c0 + 1)
            else:
                ans += c1
                res = (c0 + 1, c1)

            # print(x, ans)
            return res

        dfs(1, -1)
        return ans