3715. 完全平方祖先节点的和

Hard

题目描述

给定一个整数 n 和一个以节点 0 为根的无向树，该树有 n 个节点，编号从 0 到 n - 1。树由长度为 n - 1 的二维数组 edges 表示，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 之间有一条无向边。

同时给定一个整数数组 nums，其中 nums[i] 是分配给节点 i 的正整数。

定义值 ti 为节点 i 的祖先节点数量，使得 nums[i] * nums[ancestor] 是完全平方数。

返回范围 [1, n - 1] 内所有节点 i 的 ti 值之和。

注意：在根树中，节点 i 的祖先节点是从节点 i 到根节点 0 路径上的所有节点，不包括 i 本身。

示例 1：

输入：n = 3, edges = [[0,1],[1,2]], nums = [2,8,2]
输出：3

示例 2：

输入：n = 3, edges = [[0,1],[0,2]], nums = [1,2,4]
输出：1

示例 3：

输入：n = 4, edges = [[0,1],[0,2],[1,3]], nums = [1,2,9,4]
输出：2

约束条件：

1 <= n <= 10^5
edges.length == n - 1
edges[i] = [ui, vi]
0 <= ui, vi <= n - 1
nums.length == n
1 <= nums[i] <= 10^5
输入保证 edges 表示一棵有效的树

解题思路

这道题的核心思想是：两个数的乘积为完全平方数，当且仅当它们具有相同的"无平方因子核心"。

解题思路：

预处理无平方因子分解：对于每个数字，我们需要去除所有偶次幂的素因子，保留奇次幂的素因子。例如，36 = 2² × 3²，其无平方因子核心为 1；18 = 2 × 3²，其无平方因子核心为 2。
数学原理：两个数 a 和 b 的乘积是完全平方数，当且仅当 a 和 b 的无平方因子核心相同。这是因为完全平方数的所有素因子都是偶次幂。
算法流程：
- 使用埃拉托斯特尼筛法预计算所有数的无平方因子核心
- 构建邻接表表示树结构
- 使用 DFS 遍历树，维护祖先节点的无平方因子核心频次
- 对于每个节点，统计具有相同无平方因子核心的祖先数量
- 回溯时更新频次映射
优化要点：
- 使用频次映射避免重复计算
- DFS 过程中正确维护和回溯祖先状态
- 高效的素因子分解算法

这种方法的时间复杂度为 O(M log log M + n)，其中 M 是 nums 中的最大值，n 是节点数量。

代码实现

class Solution {
public:
    long long sumOfAncestors(int n, vector<vector<int>>& edges, vector<int>& nums) {
        int maxVal = *max_element(nums.begin(), nums.end());
        vector<int> squareFree(maxVal + 1);
        
        // 预计算无平方因子核心
        for (int i = 1; i <= maxVal; i++) {
            squareFree[i] = i;
        }
        
        for (int i = 2; i * i <= maxVal; i++) {
            if (squareFree[i] == i) {
                for (int j = i * i; j <= maxVal; j += i * i) {
                    while (squareFree[j] % (i * i) == 0) {
                        squareFree[j] /= (i * i);
                    }
                }
            }
        }
        
        // 构建邻接表
        vector<vector<int>> adj(n);
        for (auto& edge : edges) {
            adj[edge[0]].push_back(edge[1]);
            adj[edge[1]].push_back(edge[0]);
        }
        
        long long result = 0;
        unordered_map<int, int> ancestorCount;
        
        function<void(int, int)> dfs = [&](int node, int parent) {
            int sf = squareFree[nums[node]];
            
            if (node != 0) {
                result += ancestorCount[sf];
            }
            
            ancestorCount[sf]++;
            
            for (int child : adj[node]) {
                if (child != parent) {
                    dfs(child, node);
                }
            }
            
            ancestorCount[sf]--;
            if (ancestorCount[sf] == 0) {
                ancestorCount.erase(sf);
            }
        };
        
        dfs(0, -1);
        return result;
    }
};

class Solution:
    def sumOfAncestors(self, n: int, edges: List[List[int]], nums: List[int]) -> int:
        max_val = max(nums)
        square_free = list(range(max_val + 1))
        
        # 预计算无平方因子核心
        for i in range(2, int(max_val**0.5) + 1):
            if square_free[i] == i:
                square = i * i
                for j in range(square, max_val + 1, square):
                    while square_free[j] % square == 0:
                        square_free[j] //= square
        
        # 构建邻接表
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        result = 0
        ancestor_count = {}
        
        def dfs(node, parent):
            nonlocal result
            sf = square_free[nums[node]]
            
            if node != 0:
                result += ancestor_count.get(sf, 0)
            
            ancestor_count[sf] = ancestor_count.get(sf, 0) + 1
            
            for child in adj[node]:
                if child != parent:
                    dfs(child, node)
            
            ancestor_count[sf] -= 1
            if ancestor_count[sf] == 0:
                del ancestor_count[sf]
        
        dfs(0, -1)
        return result

public class Solution {
    public long SumOfAncestors(int n, int[][] edges, int[] nums) {
        int maxVal = nums.Max();
        int[] squareFree = new int[maxVal + 1];
        
        // 预计算无平方因子核心
        for (int i = 1; i <= maxVal; i++) {
            squareFree[i] = i;
        }
        
        for (int i = 2; i * i <= maxVal; i++) {
            if (squareFree[i] == i) {
                int square = i * i;
                for (int j = square; j <= maxVal; j += square) {
                    while (squareFree[j] % square == 0) {
                        squareFree[j] /= square;
                    }
                }
            }
        }
        
        // 构建邻接表
        List<int>[] adj = new List<int>[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new List<int>();
        }
        
        foreach (var edge in edges) {
            adj[edge[0]].Add(edge[1]);
            adj[edge[1]].Add(edge[0]);
        }
        
        long result = 0;
        Dictionary<int, int> ancestorCount = new Dictionary<int, int>();
        
        void Dfs(int node, int parent) {
            int sf = squareFree[nums[node]];
            
            if (node != 0) {
                ancestorCount.TryGetValue(sf, out int count);
                result += count;
            }
            
            ancestorCount.TryGetValue(sf, out int currentCount);
            ancestorCount[sf] = currentCount + 1;
            
            foreach (int child in adj[node]) {
                if (child != parent) {
                    Dfs(child, node);
                }
            }
            
            ancestorCount[sf]--;
            if (ancestorCount[sf] == 0) {
                ancestorCount.Remove(sf);
            }
        }
        
        Dfs(0, -1);
        return result;
    }
}

var sumOfAncestors = function(n, edges, nums) {
    const maxVal = Math.max(...nums);
    const squareFree = Array.from({length: maxVal + 1}, (_, i) => i);
    
    // 预计算无平方因子核心
    for (let i = 2; i * i <= maxVal; i++) {
        if (squareFree[i]

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型	复杂度
时间复杂度	O(M log log M + n)
空间复杂度	O(M + n)

其中 M 是 nums 中的最大值，n 是节点数量。时间复杂度主要由预计算无平方因子核心和 DFS 遍历组成，空间复杂度主要用于存储无平方因子数组和邻接表。

题目描述#

解题思路#

代码实现#

复杂度分析#

题目描述

解题思路

代码实现

复杂度分析