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

Question

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)

user2390182 · Accepted Answer · 2018-09-04T09:22:50.670

3

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

edited Sep 04 '18 at 09:22

answered Sep 04 '18 at 08:41

user2390182

72,016
6
67
89

Thanks, i was looking for that kiind of implementation not pasing the list as parameter. the first approach works fine for me, since i'm using python 2.7 and your second solution is only compatible for python 3.3+ – Arthur Reinhart Sep 04 '18 at 09:21
@ArthurReinhart Yup, should have mentioned that. I added a Python2 version of the generator implementation. – user2390182 Sep 04 '18 at 09:23
1

Thanks again @schwobaseggl , i haven't used generators before, and just have been learning about them cause of you sugerence. I learned a lot here https://stackoverflow.com/a/231855/8002694 – Arthur Reinhart Sep 04 '18 at 11:40

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

1 Answers1

Linked