Leetcode 2867. Count Valid Paths in a Tree

• Leetcode 2867. Count Valid Paths in a Tree
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：2867. Count Valid Paths in a Tree

1. 解题思路

这一题思路上的话由于要求路径上有且仅有一个质数点，因此，一个直接的思路就是考察所有质数的点作为中心点时，其辐射出去的非质数点的个数。

假设一个质数点$p$上有$k$个合数子节点，然后每一个合数子节点对应的只包含合数的子树当中的节点个数分别为$n_1,\cdots n_k$个，那么，包含且仅包含$p$的路径个数为：

1. $p$作为路径的一个节点：$N=\sum _{i=1}^k{n}_i$
2. $p$作为路径的一个中间节点：$N=\sum _{i=1}^k\sum _{j=1,j\ne i}^k{n}_i\times n_j=\sum _{i=1}^k{n}_i(\sum _{j=1}^k{n}_j-n_i)/2$

由此，问题就转化为只要求得每一个质数节点$p$周围相邻的合数节点$u$作为根节点时的只包含合数节点的子树的节点的个数。

而这个问题只需要用一个深度优先遍历就可以完成了。当然，为了优化执行效率，我们加上了一个cache来对其进行优化。

2. 代码实现

给出最终的python代码实现如下：

@lru_cache(None)
def get_primes():
    n = 10**5
    status = [0 for _ in range(n)]
    primes = []
    for i in range(2, n):
        if status[i] == 1:
            continue
        primes.append(i)
        for j in range(i, n, i):
            status[j] = 1
    return primes

PRIMES = get_primes()

class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        if n == 1:
            return 0
        
        primes = PRIMES[:bisect.bisect_right(PRIMES, n)]
        prime_set = set(primes)

        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
            
        @lru_cache(None)
        def dfs(u, pre):
            res = 1
            for v in graph[u]:
                if v != pre and v not in prime_set:
                    res += dfs(v, u)
            return res
        
        def query(u):
            nodes = [dfs(v, -1) for v in graph[u] if v not in prime_set]
            res = 0
            s = sum(nodes)
            for k in nodes:
                res += k * (s-k)
            return res // 2 + s
            

        return sum(query(u) for u in primes)

提交代码评测得到：耗时1683ms，占用内存71.2MB。