计算给定树中权重为完全数的节点

背景

给定一棵树，每个节点有一个权重，权重为完全数表示该节点的权重是一个正整数的平方数，如 $1, 4, 9, 16, \cdots$。请编写一个程序，计算给定树中权重为完全数的节点的个数。

思路

我们可以通过递归遍历树的每一个节点，并判断它的权重是否为完全数，来统计权重为完全数的节点个数。

判断一个数是否为完全数，可以通过判断它的平方根是否为整数来实现。

具体实现可以定义一个 isPerfectSquare(num: int) -> bool 函数用于判断给定的数字是否是平方数，然后对每个节点的权重进行判断并统计。

代码示例

下面是 Python 3 代码示例：

from typing import List
from math import isqrt

class TreeNode:
    def __init__(self, val=0, children=None):
        self.val = val
        self.children = children if children is not None else []

def isPerfectSquare(num: int) -> bool:
    root = isqrt(num)
    return root ** 2 == num

def countPerfectSquares(root: TreeNode) -> int:
    count = 0
    if root is not None:
        if isPerfectSquare(root.val):
            count += 1
        for child in root.children:
            count += countPerfectSquares(child)
    return count

使用示例

以下是一个简单的使用示例：

tree = TreeNode(3, [TreeNode(4), TreeNode(5, [TreeNode(9)])])
count = countPerfectSquares(tree)
print(count)  # 输出 2

性能分析

这个算法的时间复杂度是 $O(N)$，其中 $N$ 是树的节点数。因为每个节点都只会被访问一次，并且常数项比较小，所以实际运行效率比较高。空间复杂度同样是 $O(N)$，因为递归调用本身需要消耗一定的栈空间。