List of Depths

Given a binary tree, design an algorithm which creates a linked list of all the nodes at each depth (i.e. if you have a tree with depth D, you’ll have D linked lists).

Solution

Let’s use a breadth-first search to traverse the tree. Using the BFS, we can visit each level of the tree. We will leverage the queue to keep track of the level of the binary tree which we are visiting. When we visit a level of the tree, we will copy the queue to a temporary queue so that we can repopulate the original queue with the children of all nodes at that tree height.

We will visit all nodes in the temporary queue and do two things: append their values to a linked list which represents all nodes at that height, and append their children to the queue. In this way, we can have the algorithm visit each level of the tree, sequentially.

def create_lists(node):
  queue = []
  results = []

  queue.append(node)

  # create linked list with only the root
  linked_list = LinkedList()
  linked_list.insert(node.key)
  results.append(linked_list)

  while(queue):
    # copy queue to temporary queue
    parents = queue
    queue = []
    linked_list = LinkedList()

    for parent in parents:
      if(parent.left != None):
        linked_list.insert(parent.left.key)
        queue.append(parent.left)
      if(parent.right != None):
        linked_list.insert(parent.right.key)
        queue.append(parent.right)

    if(linked_list.head is not None):
      results.append(linked_list)

  return results

We can also implement a pre-order traversal, to visit the root node, and then all children nodes.

def pre_order_traversal(node, results, level):
  if(node == None):
    return

  if(len(results) == level):
    linked_list = LinkedList()
    linked_list.insert(node.key)
    results.append(linked_list)
  else:
    results[level].insert(node.key)

  pre_order_traversal(node.left, results, level+1)
  pre_order_traversal(node.right, results, level+1)