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

链接：2867.统计树中的合法路径数目
难度：Hard
标签：树、深度优先搜索、数学、动态规划、数论
简介：请你返回树中的 合法路径数目 。

题解 1 - python

• 编辑时间：2024-02-27
• 执行用时：322ms
• 内存消耗：58.63MB
• 编程语言：python
• 解法介绍：预处理好质数表，通过遍历所有质数，找到以当前质数为根结点的时候，所有子树的长度，进行两两相乘。

def get_primes2(n: int) -> List[bool]:
        n += 1
        primes = [True for _ in range(n)]
        primes[0] = primes[1] = False
        for i in range(2, n):
            if primes[i]:
                j = 2
                while i * j < n:
                    primes[i*j] = False
                    j += 1
        return primes
    primes = get_primes2(10 ** 5 + 1)
    
    class Solution:
        def countPaths(self, n: int, edges: List[List[int]]) -> int:
            nodes = [[] for _ in range(n + 1)]
            for n1, n2 in edges:
                nodes[n1].append(n2)
                nodes[n2].append(n1)
            ans = 0
    
            cache = defaultdict(int)
            def dfs(arr: List[int], node: int, parent: int):
                if primes[node]: return
                arr.append(node)
                ans = 1
                for child in nodes[node]:
                    if not primes[child] and child != parent:
                        ans += dfs(arr, child, node)
                return ans
    
            for node in range(1, n + 1):
                if primes[node]:
                    cur = 0
                    for child in nodes[node]:
                        if not primes[child]:
                            if child not in cache:
                                arr = []
                                res = dfs(arr, child, node)
                                for item in arr: cache[item] = res
                            ans += cache[child] * cur
                            cur += cache[child]
                    ans += cur
        return ans