binary_tree_level_order_traversal

Solution

	
	from collections import deque
from typing import List, Optional

from publicmatt.leet import TreeNode


def level_order(root: Optional[TreeNode]) -> List[List[int]]:
    out = []

    q = deque()
    q.append(root)
    while q:
        size = len(q)
        nodes = []
        for i in range(size):
            node = q.popleft()
            if node:
                nodes.append(node.val)
                q.append(node.left)
                q.append(node.right)
        if nodes:
            out.append(nodes)
    return out

	

Tests

	
	from publicmatt.leet import TreeNode
from publicmatt.leet.problems.binary_tree_level_order_traversal import level_order


def test_simple_level():
    arr = [3, 9, 20, None, None, 15, 7]
    tree = TreeNode.from_list(arr)
    assert level_order(tree) == [[3], [9, 20], [15, 7]]


def test_single_level():
    arr = [1]
    tree = TreeNode.from_list(arr)
    assert level_order(tree) == [[1]]


def test_empty():
    arr = []
    tree = TreeNode.from_list(arr)
    assert level_order(tree) == []