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