Chapter 4 Trees and Graphs - 4.4

Problem 4.4: Given a binary search tree, design an algorithm which creates a linked list of all the nodes each depth (i.e., if you have a tree with depthD, you'll have D linked lists).

BFS, again...

from min_binary_tree import *
from queue import *
def btree_to_lists(root):
    lists = []
    lists.append([])
    # Use BFS to traverse by level order
    q = queue()
    q.enqueue(root)
    current_level = 0
    current_level_num = 1
    while not q.is_empty():
        n = q.dequeue()
        lists[current_level].append(n.value)
        current_level_num -= 1
        if n.left != None:
            q.enqueue(n.left)
        if n.right != None:
            q.enqueue(n.right)
        # End of a level
        if current_level_num == 0:
            lists.append([])
            current_level += 1
            current_level_num = 2**current_level
    return lists
# Test case
if __name__ == "__main__":
    array = [i for i in range(0, 8)]
    root = construct_min_btree(array, 0, 7)
    print btree_to_lists(root)