Skip to content

Sort Binary Tree by Levels using Python

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
Code language: Python (python)

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:

Code language: JSON / JSON with Comments (json)

Example 2 – following tree:

1 8 4 3 5 7

Should return following list:

Code language: JSON / JSON with Comments (json)

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])
Code language: Python (python)

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 []
Code language: Python (python)

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 []
Code language: plaintext (plaintext)

See also  How to Create Your Own Python Split Function
Notify of
1 Comment
Oldest Most Voted
Inline Feedbacks
View all comments
2 months ago

Line 7 of the the second option should be v = q.pop(0)with a zero. Otherwise it will append the branches from right to left.

Would love your thoughts, please comment.x