Without testing, as I don't have your Tree
class, but I hope the general idea applies:
def levelorder(tree):
levels = [ [tree.root] ]
while levels [-1]:
nextLevel = []
for node in levels [-1]:
nextLevel.extend ( [node for node in (node.left, node.right) if node] )
levels.append (nextLevel)
return levels # or maybe levels [:-1]
Fully working example:
#! /usr/bin/python3
class Node:
def __init__(self, payload, left = None, right = None):
self.payload = payload
self.left = left
self.right = right
def __repr__(self):
return self.payload
class Tree:
def __init__(self, root):
self.root = root
def levels(self):
levels = [[self.root]]
while levels[-1]:
nextLevel = []
for node in levels[-1]:
nextLevel.extend([node for node in (node.left, node.right) if node])
levels.append(nextLevel)
return levels[:-1]
t = Tree(Node('41', Node('7', Node('1'), Node('19')), Node('53', Node('47'))))
print(t.levels())
Output is [[41], [7, 53], [1, 19, 47]]