Leetcode 3715. Sum of Perfect Square Ancestors

• [Leetcode 3715. Sum of Perfect Square Ancestors](#Leetcode 3715. Sum of Perfect Square Ancestors)
  • [1. 解题思路](#1. 解题思路)
  • [2. 代码实现](#2. 代码实现)

• 题目链接：3715. Sum of Perfect Square Ancestors

1. 解题思路

这一题整体的思路上就是一个树的遍历，我们遍历每一个节点，然后考察其所有的父节点当中有多少个元素与之相乘可以得到完全平方数即可。

但是，如果要暴力相乘然后判断平方数事实上还是非常麻烦的，因此这里我们做了一下优化，即对于每一个数，我们可以事先将其所有的平方数全部除掉，此时，找到平方数的操作就可以简化为求出其所有的父节点当中有多少与当前元素相同的元素即可。

2. 代码实现

给出python代码实现如下：

def get_perfect_square(n):
    ans = []
    for i in range(2, n):
        if i * i > n:
            break
        ans.append(i*i)
    return ans

PerfectSquares = get_perfect_square(10**5)

class Solution:
    def sumOfAncestors(self, n: int, edges: List[List[int]], nums: List[int]) -> int:
        
        @lru_cache(None)
        def simplify(num):
            for ps in PerfectSquares:
                if ps > num:
                    break
                while num % ps == 0:
                    num = num // ps
            return num
        
        nums = [simplify(num) for num in nums]
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        ans = 0
        def dfs(u, p, history):
            nonlocal ans
            ans += history[nums[u]]
            history[nums[u]] += 1
            for v in graph[u]:
                if v == p:
                    continue
                dfs(v, u, history)
            history[nums[u]] -= 1
            return

        dfs(0, -1, defaultdict(int))
        return ans

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