Kimi LeetCode 3715. 完全平方数的祖先个数总和 Python3实现

以下是 LeetCode 3715. 完全平方数的祖先个数总和 的 Python3 实现基于 DFS 哈希表 无平方因子核预处理的思路。---核心思路关键数学性质两个数 a × b 是完全平方数当且仅当它们去掉所有平方因子后的剩余部分无平方因子核相等。例如- 8 2³去掉 2²偶数指数剩余 2¹ → core[8] 2- 2 2¹ → core[2] 2- 16 2⁴ → 所有指数都是偶数 → core[16] 1所以 nums[i] × nums[ancestor] 为完全平方数 ⇔ core[nums[i]] core[nums[ancestor]]算法1. 预处理所有数的无平方因子核 core[x]2. DFS 遍历树用哈希表 cnt 记录当前路径上各 core 值的出现次数3. 对每个非根节点查询 cnt[core[nums[i]]] 即为符合条件的祖先个数4. 回溯时恢复现场---Python3 代码pythonfrom typing import Listfrom collections import defaultdictclass Solution:# 预处理计算每个数的无平方因子核square-free core# core[x] 去掉 x 中所有平方因子后剩余的部分MX 100001core [0] * MXclassmethoddef _init_core(cls):类方法只需执行一次预处理if cls.core[1] ! 0:return # 已经初始化过for i in range(1, cls.MX):if cls.core[i] 0: # i 本身是无平方因子数j 1while i * j * j cls.MX:cls.core[i * j * j] ij 1def sumOfAncestors(self, n: int, edges: List[List[int]], nums: List[int]) - int:self._init_core()# 建图无向树g [[] for _ in range(n)]for u, v in edges:g[u].append(v)g[v].append(u)ans 0# cnt[c] 当前DFS路径上 core 值为 c 的节点个数cnt defaultdict(int)def dfs(x: int, fa: int):nonlocal ansc self.core[nums[x]]# 当前节点的祖先中core 值相同的个数ans cnt[c]cnt[c] 1for y in g[x]:if y ! fa:dfs(y, x)# 回溯恢复现场cnt[c] - 1dfs(0, -1)return ans---复杂度分析项目 复杂度预处理 core O(MX \cdot \sqrt{MX})实际约 10^5 \times 10建图 O(n)DFS 遍历 O(n)总时间 O(n MX \cdot \sqrt{MX})空间 O(n MX)---示例验证示例 1n3, edges[[0,1],[1,2]], nums[2,8,2]节点 nums core 祖先core计数 贡献0 2 2 {} 01 8 2 {2:1} 12 2 2 {2:2} 2总和 0 1 2 3 ✓示例 2n3, edges[[0,1],[0,2]], nums[1,2,4]节点 nums core 祖先core计数 贡献0 1 1 {} 01 2 2 {1:1} 02 4 1 {1:1} 1总和 0 0 1 1 ✓---另一种写法更简洁直接修改 nums如果不想用类变量预处理也可以直接对每个 nums[i] 做质因数分解去掉所有平方因子pythonfrom typing import Listfrom collections import defaultdictclass Solution:def sumOfAncestors(self, n: int, edges: List[List[int]], nums: List[int]) - int:# 预处理去掉每个 nums[i] 中的所有平方因子for i in range(n):x nums[i]d 2while d * d x:while x % (d * d) 0:x // d * dd 1nums[i] xg [[] for _ in range(n)]for u, v in edges:g[u].append(v)g[v].append(u)ans 0cnt defaultdict(int)def dfs(x: int, fa: int):nonlocal ansc nums[x]ans cnt[c]cnt[c] 1for y in g[x]:if y ! fa:dfs(y, x)cnt[c] - 1dfs(0, -1)return ans两种方法本质相同第二种更简洁但第一种预处理 core 数组在多次调用时效率更高。两种方法都能通过 n \le 10^5 的数据范围。