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

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

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/

For reference

# Python3 Program to print BFS traversal
# from a given source vertex. BFS(int s)
# traverses vertices reachable from s.
from collections import defaultdict
 
# This class represents a directed graph
# using adjacency list representation
class Graph:
 
    # Constructor
    def __init__(self):
 
        # default dictionary to store graph
        self.graph = defaultdict(list)
 
    # function to add an edge to graph
    def addEdge(self,u,v):
        self.graph[u].append(v)
 
    # Function to print a BFS of graph
    def BFS(self, s):
 
        # Mark all the vertices as not visited
        visited = [False] * (max(self.graph) + 1)
 
        # Create a queue for BFS
        queue = []
 
        # Mark the source node as
        # visited and enqueue it
        queue.append(s)
        visited[s] = True
 
        while queue:
 
            # Dequeue a vertex from
            # queue and print it
            s = queue.pop(0)
            print (s, end = " ")
 
            # Get all adjacent vertices of the
            # dequeued vertex s. If a adjacent
            # has not been visited, then mark it
            # visited and enqueue it
            for i in self.graph[s]:
                if visited[i] == False:
                    queue.append(i)
                    visited[i] = True
 
# Driver code
 
# Create a graph given in
# the above diagram
g = Graph()
g.addEdge(0, 1)
g.addEdge(0, 2)
g.addEdge(1, 2)
g.addEdge(2, 0)
g.addEdge(2, 3)
g.addEdge(3, 3)
 
print ("Following is Breadth First Traversal"
                  " (starting from vertex 2)")
g.BFS(2)
 
# This code is contributed by Neelam Yadav

Answer 3

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