python - 递归计算二叉树中的节点

Question

我必须递归地计算二叉树中的节点。我是 python 的新手，找不到任何解决我的问题的方法来完成我的代码。

这是我已经尝试过的。如您所见，它还不完整，我不知道该去哪里。

class Tree:
    def __init__(self, root):
        self.root = root

    def add(self, subtree):
        self.root.children.append(subtree)

class Node:
    def __init__(self, value, children=None):
        self.value = value
        self.children = children if children is not None else []

def check_root(tree):
    if tree.root is None:
        return 0
    if tree.root is not None:
        return count_nodes(tree)

def count_nodes(tree):
    if tree.root.children is not None:
        j = 0
        for i in tree.root.children:
            j = 1 + count_nodes(tree)
        return j


print(count_nodes(Tree(None))) # 0
print(count_nodes(Tree(Node(10)))) # 1
print(count_nodes(Tree(Node(5, [Node(6), Node(17)])))) #3

随着每一个新步骤，我都会遇到不同的错误。例如。使用这段代码，我已经超过了最大递归深度。

感谢您花时间阅读本文。任何提示或帮助下一步做什么将不胜感激。

Answer 1

我首先将根节点传递给 count_nodes 函数 -

print(count_nodes(Tree(None)).root) # 0
print(count_nodes(Tree(Node(10))).root) # 1
print(count_nodes(Tree(Node(5, [Node(6), Node(17)]))).root) #3

或者为此创建一个辅助函数。

然后 count_nodes 函数可以看起来像这样

def count_nodes(node):
    return 1 + sum(count_nodes(child) for child in node.children)

编辑:我刚刚注意到，您可以有一个 None 根，这意味着您还应该处理它:

def count_nodes(node):
    if node is None:
        return 0
    return 1 + sum(count_nodes(child) for child in node.children)

如果你真的想在一个函数中处理树或节点，你可以让它更丑一点:

def count_nodes(tree_or_node):
    if isinstance(tree_or_node, Tree):
        return count_nodes(tree_or_node.root)
    if tree_or_node is None:
        return 0
    return 1 + sum(count_nodes(child) for child in tree_or_node.children)

然后你就可以像原来那样调用它了。