用 Python 获取二叉树中给定层级上的所有节点

3
我正在尝试获取Python二叉树中节点(对象)的列表。我正在寻找在节点对象上实现的递归函数,这样我就可以在根节点上调用该函数,并将其沿着子节点向下传递直到达到特定级别,然后它将以列表形式返回这些节点。
目前我的方法,我不确定这是否是正确或最佳的实现方式:
def get_level_nodes(self, nodes, level=1):
    if self.level > level:
        return nodes
    if self.level == level:
        nodes.append(self)
        return nodes
    for child in self.child_id:
        nodes += child.get_level_nodes(node, level)
    return nodes

# Getting the list
nodes_list = root_node.get_level_nodes([], 3)
1个回答

3

没有必要在节点间传递节点列表。每个节点只需返回其自身子树中相应级别的节点,并将邻居的组合留给父节点处理:

def get_level_nodes(self, level=1):
    if self.level > level:
        return []
    if self.level == level:
        return [self]
    # child_id seems an odd name
    return [n for c in self.children for n in c.get_level_nodes(level)]

一个更节省空间的实现方式是使用生成器函数,它不需要为每个子树构建中间列表:
def get_level_nodes(self, level=1):
    if self.level > level:
        return
    if self.level == level:
        yield self
    else:
        for c in self.children:
            for n in c.get_level_nodes(level):
                yield n
            # or in Python3
            # yield from c.get_level_nodes(level)


nodes_list = list(root_node.get_level_nodes(3))

谢谢,我正在寻找这种实现方式而不是将列表作为参数传递。第一种方法对我来说很好,因为我正在使用Python 2.7,而您的第二种解决方案只适用于Python 3.3+。 - Arthur Reinhart
@ArthurReinhart 是的,我应该提到这一点。我添加了一个生成器实现的Python2版本。 - user2390182
1
再次感谢 @schwobaseggl,因为您的建议,我开始了解生成器,之前从未使用过。我在这里学到了很多 https://dev59.com/yXVC5IYBdhLWcg3woSpW#231855 - Arthur Reinhart

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接