Binary tree get all nodes for every level

Question

I am trying to solve the following, say I have this binary tree...

       3
      / \
     9  20
       /  \
      15   7

Then I need to get all nodes at every level, so my result will be...

[[3],[9,20],[15,7]]

I believe I am getting close to a solution, but I'm not sure where I'm going wrong, if someone can help me with my solution that would be great, if I am on the wrong track, please let me know.

My first step was to get the depth of the tree using the following function...

def get_depth(self, root):
    if root.left == None and root.right == None:
        return 1
    return max(self.get_depth(root.left), self.get_depth(root.right)) + 1

The depth is 3 .

Next I called a function that was intended to give me the expected result...

def levelOrder(self, root):
    depth = self.get_depth(root)
    return self.find_comb(root, [[]]*depth, 0)

def find_comb(self, root, result, level):
    if root is None:
        return None
    self.find_comb(root.left, result, level+1)
    self.find_comb(root.right, result, level+1)
    result[level].append(root.val)
    return result

My thought process was, I would recurse through the tree and the level parameter would keep track of the current level that I am on. Then I would append all root.val on that level to the result at that index.

So let's say I am on level 1 (we start at 0), then the nodes at level 1 are 9 and 20 . So the result would look like [[], [9, 20], [], []] , and it would do this for every level of the tree. I hope my logic is clear. Instead the result I am getting is...

[[9, 15, 7, 20, 3], [9, 15, 7, 20, 3], [9, 15, 7, 20, 3]]

Answer 1

You actually don't need to find the depth of the tree. You just traverse the tree, while keeping the level at which are you, let's say in level variable. And insert your value at listOfLists[level] . Of course you have to be handle the IndexError: list index out of range . Here's a simple implementation:

def levelTraversal(root, level):
    if root is None:
        return

    if level >= len(listOfLists):
        list = []
        listOfLists.append(list)
    listOfLists[level].append(root.v)
    levelTraversal(root.l, level+1)
    levelTraversal(root.r, level+1)

Call it like this : levelTraversal(rootNode, 0)

For your tree the result is : [[3], [9, 20], [15, 7]]

Answer 2

I don't know Phyton but you might want to look at a BFS implementation on Phyton.

https://www.geeksforgeeks.org/breadth-first-traversal-for-a-graph/

Answer 3

Few improvements and variations from answer of @Schidu Luca

The problem is the way it was implemented, the child node loses which parent is despite we can retrieve at which level of the tree is.

Add mapping child -> parent

nodes = []

def traverse(root, level, parent):
    if not root:
        return
    if level >= len(nodes):
        nodes_level = []
        nodes.append(nodes_level)
    parent_key = parent and parent.value or None
    nodes[level].append({root.value: parent_key})
    traverse(root.left, level + 1, root)
    traverse(root.right, level + 1, root)


traverse(self.root, 0, None)
return nodes

Mapping level -> parent -> childs

nodes = {}
def traverse_mapped(root, level, parent):
    if not root:
        return
    if level >= len(nodes):
        nodes_level = {}
        nodes[level] = nodes_level
    parent_key = parent and parent.value or None
    if parent_key in nodes[level]:
        nodes[level][parent_key].append(root.value)
    else:
        nodes[level][parent_key] = [root.value]
    traverse_mapped(root.left, level + 1, root)
    traverse_mapped(root.right, level + 1, root)

traverse_mapped(self.root, 0, None)
return nodes