2026-03-09：完全平方数的祖先个数总和。用go语言，给出一个节点数为 n（编号 0 到 n−1）的无向树，并把节点 0 视为根。树中的边由长度

福大大架构师每日一题

发布于 2026-03-09 13:59:58

740

2026-03-09：完全平方数的祖先个数总和。用go语言，给出一个节点数为 n（编号 0 到 n−1）的无向树，并把节点 0 视为根。树中的边由长度为 n−1 的数组 edges 给出，edges[i] = [u_i, v_i] 表示节点 u_i 和 v_i 之间有一条边。另有一个正整数数组 nums，其中 nums[i] 表示分配给节点 i 的值。把树当成有根树来看：对于任意节点 i（i ≠ 0），其“祖先”指的是从 i 沿着通往根 0 的路径上、除去自身以外出现的那些节点。定义 t_i 为在这些祖先中，与节点 i 相乘后能得到完全平方数的祖先数量（完全平方数指某个整数的平方，如 1、4、9 等）。

任务是计算并返回所有节点 i = 1,2,…,n−1 的 t_i 之和。

1 <= n <= 100000。

edges.length == n - 1。

edges[i] = [ui, vi]。

0 <= ui, vi <= n - 1。

nums.length == n。

1 <= nums[i] <= 100000。

输入保证 edges 表示一棵有效的树。

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

输出： 3。

解释：

i	祖先	nums[i] * nums[ancestor]	平方数检查	ti
1	[0]	nums[1] * nums[0] = 8 * 2 = 16	16 是完全平方数	1
2	[1, 0]	nums[2] * nums[1] = 2 * 8 = 16nums[2] * nums[0] = 2 * 2 = 4	16 和 4 都是完全平方数	2

因此，所有非根节点的有效祖先配对总数为 1 + 2 = 3。

题目来自力扣3715。

一、关键前置知识：完全平方数的数学性质

两个数相乘是完全平方数的充要条件：将两个数分解质因数后，所有质因数的指数均为偶数。简化方法：对任意数 x，去掉其质因数分解中所有指数为偶数的部分（只保留指数为奇数的质因数，且指数变为1），得到的数称为 x 的“无平方因子核”（代码中记为 core[x]）。例如：

• 8 = 2³ → 去掉2²（偶数指数），剩余2¹ → core[8] = 2
• 2 = 2¹ → core[2] = 2
• 16 = 2⁴ → 所有指数都是偶数 → core[16] = 1 此时，a × b 是完全平方数 ⇨ core[a] = core[b]。

二、代码执行的完整过程（分步骤）

步骤1：预处理生成所有数的“无平方因子核”（init函数）

代码首先定义了一个长度为100001的数组 core，并通过init函数预处理1~100000所有数的核：

1. 初始化 core 数组所有元素为0；
2. 遍历每个数 i（从1到100000）：
- • 如果 core[i] 为0（说明i是未处理的“核数”，即自身无平方因子）；
- • 遍历 j 从1开始，计算 i × j²（给i乘任意完全平方数），将这些数的core值设为i；
- • 例如：i=2时，j=1 → 2×1=2（core[2]=2），j=2 → 2×4=8（core[8]=2），j=3 → 2×9=18（core[18]=2），直到i×j² ≥ 100001停止。最终，core[x] 即为x的无平方因子核，后续只需比较两个数的core是否相等，就能判断乘积是否为完全平方数。

步骤2：构建树的邻接表

输入的树是无向边（如[0,1]），需要转换为邻接表g（二维数组）：

• g[x] 存储所有与x直接相连的节点；
• 例如输入edges = [[0,1],[1,2]]，则g[0] = [1]，g[1] = [0,2]，g[2] = [1]。

步骤3：深度优先遍历（DFS）统计符合条件的祖先数量

核心逻辑是回溯式DFS：从根节点0出发，遍历每个子节点，记录路径上（祖先）的core值出现次数，统计当前节点的符合条件祖先数，遍历完子树后恢复现场（避免影响其他分支）。以输入n=3, nums=[2,8,2]为例，详细执行过程：

子步骤3.1：初始化计数哈希表

定义cnt（map类型），用于记录当前DFS路径上（即当前节点的所有祖先）各core值的出现次数，初始为空。

子步骤3.2：执行DFS（起点：根节点0，父节点-1）

1. 处理节点0（根节点）：
- • 计算core[nums[0]] = core[2] = 2；
- • 根节点无祖先，因此ans（总和）无增加；
- • 将cnt[2]加1（此时cnt = {2:1}）；
- • 遍历邻接节点：只有1（父节点是-1，跳过父节点判断），递归处理节点1。
2. 处理节点1（父节点0）：
- • 计算core[nums[1]] = core[8] = 2；
- • 统计祖先中core=2的数量：cnt[2] = 1，因此ans += 1（此时ans=1）；
- • 将cnt[2]加1（此时cnt = {2:2}）；
- • 遍历邻接节点：0（父节点，跳过）、2，递归处理节点2。
3. 处理节点2（父节点1）：
- • 计算core[nums[2]] = core[2] = 2；
- • 统计祖先中core=2的数量：cnt[2] = 2，因此ans += 2（此时ans=3）；
- • 将cnt[2]加1（此时cnt = {2:3}）；
- • 遍历邻接节点：1（父节点，跳过），无其他节点，递归返回；
- • 恢复现场：cnt[2]减1（回到cnt = {2:2}）。
4. 节点1处理完毕，恢复现场：
- • cnt[2]减1（回到cnt = {2:1}），递归返回。
5. 节点0处理完毕，恢复现场：
- • cnt[2]减1（回到cnt = {}），DFS结束。

最终ans=3，与题目输出一致。

步骤4：返回结果

DFS结束后，ans即为所有非根节点的t_i之和，返回该值。

三、时间复杂度分析

1. 预处理阶段（init函数）

• 遍历i从1到100000，对每个i，遍历j直到i×j² ≥ 100001；
• 时间复杂度：O(M)（M=100000），因为所有i×j²的总操作数等价于对每个数x只处理一次，总次数约为M + M/4 + M/9 + ... ≈ M × π²/6 ≈ 1.64M，属于O(M)级别。

2. 构建邻接表

• 遍历n-1条边，每条边添加两个方向的邻接关系，时间复杂度O(n)。

3. DFS遍历

• 每个节点仅被访问一次，每条边仅被处理两次（无向边），时间复杂度O(n)；
• 哈希表cnt的操作（增/减/查）均为O(1)（平均情况）。

总时间复杂度

O(M + n)，其中M=100000（固定值），n≤100000，因此可简化为O(10^5)（或O(max(M, n))）。

四、额外空间复杂度分析

1. 预处理阶段

• core数组长度为100001，空间复杂度O(M)（M=100000）。

2. 邻接表

• 存储n个节点的邻接关系，总边数为2(n-1)，空间复杂度O(n)。

3. DFS相关

• 哈希表cnt：最坏情况下（树退化为链表），存储路径上所有节点的core值，最多n个键值对，空间复杂度O(n)；
• DFS递归栈：最坏情况下递归深度为n（链表），空间复杂度O(n)。

总额外空间复杂度

O(M + n)，同样可简化为O(10^5)（或O(max(M, n))）。

总结

1. 核心逻辑：利用“无平方因子核”将“乘积为完全平方数”转化为“核相等”，通过预处理避免重复计算，再用DFS回溯统计路径上的核出现次数；
2. 时间复杂度：O(max(10^5, n))（预处理10^5次，DFSn次）；
3. 空间复杂度：O(max(10^5, n))（core数组占10^5空间，邻接表/DFS栈/哈希表占n空间）。

Go完整代码如下：

package main

import (
    "fmt"
)

// 预处理 core
const mx = 100_001

var core = [mx]int{}

func init() {
    for i := 1; i < mx; i++ {
        if core[i] == 0 {
            for j := 1; i*j*j < mx; j++ {
                core[i*j*j] = i
            }
        }
    }
}

func sumOfAncestors(n int, edges [][]int, nums []int) (ans int64) {
    g := make([][]int, n)
    for _, e := range edges {
        x, y := e[0], e[1]
        g[x] = append(g[x], y)
        g[y] = append(g[y], x)
    }

    cnt := map[int]int{}
    var dfs func(int, int)
    dfs = func(x, fa int) {
        c := core[nums[x]]
        // 本题 x 的祖先不包括 x 自己
        ans += int64(cnt[c])
        cnt[c]++
        for _, y := range g[x] {
            if y != fa {
                dfs(y, x)
            }
        }
        cnt[c]-- // 恢复现场
    }
    dfs(0, -1)
    return
}

func main() {
    n := 3
    edges := [][]int{{0, 1}, {1, 2}}
    nums := []int{2, 8, 2}
    result := sumOfAncestors(n, edges, nums)
    fmt.Println(result)
}

Python完整代码如下：

# -*-coding:utf-8-*-

import sys
from typing import List, Dict, Set

# 设置递归深度
sys.setrecursionlimit(200000)

# 预处理 core 数组
mx = 100_001
core = [0] * mx

# 初始化 core 数组
for i in range(1, mx):
    if core[i] == 0:
        j = 1
        while i * j * j < mx:
            core[i * j * j] = i
            j += 1

def sumOfAncestors(n: int, edges: List[List[int]], nums: List[int]) -> int:
    # 构建邻接表
    g = [[] for _ in range(n)]
    for x, y in edges:
        g[x].append(y)
        g[y].append(x)
    
    ans = 0
    cnt: Dict[int, int] = {}
    
    def dfs(x: int, fa: int) -> None:
        nonlocal ans
        c = core[nums[x]]
        # 本题 x 的祖先不包括 x 自己
        ans += cnt.get(c, 0)
        cnt[c] = cnt.get(c, 0) + 1
        
        for y in g[x]:
            if y != fa:
                dfs(y, x)
        
        cnt[c] -= 1
        if cnt[c] == 0:
            del cnt[c]
    
    dfs(0, -1)
    return ans

def main():
    n = 3
    edges = [[0, 1], [1, 2]]
    nums = [2, 8, 2]
    result = sumOfAncestors(n, edges, nums)
    print(result)

if __name__ == "__main__":
    main()

C++完整代码如下：

#include <iostream>
#include <vector>
#include <unordered_map>
#include <functional>

using namespace std;

// 预处理 core
constint mx = 100001;
int core[mx] = {0};

// 使用静态初始化函数
static int init_core() {
    for (int i = 1; i < mx; i++) {
        if (core[i] == 0) {
            for (int j = 1; i * j * j < mx; j++) {
                core[i * j * j] = i;
            }
        }
    }
    return0;
}

// 在程序启动时执行初始化
static int _init = init_core();

long long sumOfAncestors(int n, vector<vector<int>>& edges, vector<int>& nums) {
    // 构建邻接表
    vector<vector<int>> g(n);
    for (auto& e : edges) {
        int x = e[0], y = e[1];
        g[x].push_back(y);
        g[y].push_back(x);
    }

    long long ans = 0;
    unordered_map<int, int> cnt;

    // 定义DFS函数
    function<void(int, int)> dfs = [&](int x, int fa) {
        int c = core[nums[x]];
        // 本题 x 的祖先不包括 x 自己
        ans += cnt[c];
        cnt[c]++;

        for (int y : g[x]) {
            if (y != fa) {
                dfs(y, x);
            }
        }

        // 恢复现场
        cnt[c]--;
        if (cnt[c] == 0) {
            cnt.erase(c);
        }
    };

    dfs(0, -1);
    return ans;
}

int main() {
    int n = 3;
    vector<vector<int>> edges = {{0, 1}, {1, 2}};
    vector<int> nums = {2, 8, 2};
    long long result = sumOfAncestors(n, edges, nums);
    cout << result << endl;

    return0;
}