CODEWITHSUNDEEP
pythontreebinary
Arthur Reinhart
Arthur Reinhart

Reputation: 63

Get all nodes on a given level in a binary tree with python

i am trying to get a list of nodes (objetcs) in a python binary tree, i am looking for a recursive function implemented in the node object, so i will call function on the root node, and it will going down on childs nodes till reachs the specific level, and then will return those nodes in a list

My current aproach, i'm not sure if this is correct or the best way to implement it:

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)

Upvotes: 3

Views: 3491

Answers (1)

user2390182
user2390182

Reputation: 73450

There is no real need to pass a list of nodes around. Each node can just return the appropriate level-nodes of its own subtree, and leave the combining of neighbours to the parent:

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)]

A more space-efficient implementation that does not build intermediate lists for each subtree would be a generator function:

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))

Upvotes: 3

Related Questions