The challenge
You are given a binary tree:
class Node:
def __init__(self, L, R, n):
self.left = L
self.right = R
self.value = n
Your task is to return the list with elements from tree sorted by levels, which means the root element goes first, then root children (from left to right) are second and third, and so on.
Return empty list if root is None.
Example 1 – following tree:
2
8 9
1 3 4 5
Should return following list:
[2,8,9,1,3,4,5]
Example 2 – following tree:
1
8 4
3 5
7
Should return following list:
[1,8,4,3,5,7]
Test cases
Test.assert_equals(tree_by_levels(None), [])
Test.assert_equals(tree_by_levels(Node(Node(None, Node(None, None, 4), 2), Node(Node(None, None, 5), Node(None, None, 6), 3), 1)), [1, 2, 3, 4, 5, 6])
The solution in Python
# Our helper function
# ..takes in `node`
def tree_iterator(node: Node):
# create a list to loop through
nodes = [node]
# loop
while nodes:
# yield from the list
yield from nodes
# internal loop
for n in nodes[:]:
# add to list if exists
if n.left: nodes.append(n.left)
if n.right: nodes.append(n.right)
# remove from teh main list loop
nodes.remove(n)
# The primary function being called
# ..passes in `node`
def tree_by_levels(node):
# return a list of values from our helper function
# otherwise return `[]` if node is empty
return [n.value for n in tree_iterator(node)] if node else []
Another option:
def tree_by_levels(node):
# create a return list, and a queue to loop
p, q = [], [node]
# loop
while q:
# take the first item from the queue
v = q.pop(0)
# if it is not empty
if v is not None:
# add it's value to the return list
p.append(v.value)
# add the left and right nodes to the queue
q += [v.left,v.right]
# return the final list, otherwise return [] is empty
return p if not node is None else []