See all coding puzzles and patterns here.

Problem Definition

Given a binary tree, return the vertical order traversal of its nodes - i.e., the nodes at each level from left to right.

Examples

# Tree:
#      1
#    /   \
#   2     3
#  / \   / \
# 4   5 6   7
#    /   \   \
#   8     9   10

eight = Node(8)
ten = Node(10)
seven = Node(7, right=ten)
nine = Node(9)
six = Node(6, right=nine)
five = Node(5, left=eight)
four = Node(4)
three = Node(3, left=six, right=seven)
two = Node(2, left=four, right=five)
one = Node(1, left=two, right=three)

one_result = [
    [1],
    [2, 3],
    [4, 5, 6, 7],
    [8, 9, 10],
]
two_result = [[2], [4, 5], [8]]
three_result = [[3], [6, 7], [9, 10]]

assert order_traversal(one) == one_result
assert order_traversal(two) == two_result
assert order_traversal(three) == three_result

Solution

We implement a BFS solution. 

from collections import deque

def order_traversal(root) -> list[list[int]]:
    q = deque([root])
    ans = []
    while q:
        curr = []
        for _ in range(len(q)):
            item = q.popleft()
            curr.append(item.val)
            if item.left:
                q.append(item.left)
            if item.right:
                q.append(item.right)
        ans.append(curr)
    return ans

Complexity

Let $n = $ the number of nodes in the tree.

Because we're visiting every node once, runtime complexity should be $\Theta(n)$.

Since we're storing a copy of the value for each node, the auxiliary space complexity should be $O(n)$.