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

发布时间：2026/7/27 6:43:44

以下是 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 代码

```python
from typing import List
from collections import defaultdict

class Solution:
# 预处理：计算每个数的无平方因子核（square-free core）
# core[x] = 去掉 x 中所有平方因子后剩余的部分
MX = 100001
core = [0] * MX

@classmethod
def _init_core(cls):
"""类方法，只需执行一次预处理"""
if cls.core[1] != 0:
return # 已经初始化过
for i in range(1, cls.MX):
if cls.core[i] == 0: # i 本身是无平方因子数
j = 1
while i * j * j < cls.MX:
cls.core[i * j * j] = i
j += 1

def 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 ans
c = self.core[nums[x]]
# 当前节点的祖先中，core 值相同的个数
ans += cnt[c]
cnt[c] += 1

for y in g[x]:
if y != fa:
dfs(y, x)

# 回溯：恢复现场
cnt[c] -= 1

dfs(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)

---

示例验证

示例 1：`n=3, edges=[[0,1],[1,2]], nums=[2,8,2]`

节点 nums core 祖先core计数贡献
0 2 2 `{}` 0
1 8 2 `{2:1}` 1
2 2 2 `{2:2}` 2

总和 = `0 + 1 + 2 = 3` ✓

示例 2：`n=3, edges=[[0,1],[0,2]], nums=[1,2,4]`

节点 nums core 祖先core计数贡献
0 1 1 `{}` 0
1 2 2 `{1:1}` 0
2 4 1 `{1:1}` 1

总和 = `0 + 0 + 1 = 1` ✓

---

另一种写法（更简洁，直接修改 nums）

如果不想用类变量预处理，也可以直接对每个 `nums[i]` 做质因数分解，去掉所有平方因子：

```python
from typing import List
from collections import defaultdict

class 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 = 2
while d * d <= x:
while x % (d * d) == 0:
x //= d * d
d += 1
nums[i] = x

g = [[] for _ in range(n)]
for u, v in edges:
g[u].append(v)
g[v].append(u)

ans = 0
cnt = defaultdict(int)

def dfs(x: int, fa: int):
nonlocal ans
c = nums[x]
ans += cnt[c]
cnt[c] += 1
for y in g[x]:
if y != fa:
dfs(y, x)
cnt[c] -= 1

dfs(0, -1)
return ans
```