Python implementations of some common binary tree operations, such as breadth-first traversal (BFS) and depth-first traversals (DFS; pre-order, in-order, post-order, both recursive and iterative).

binary_trees.py

#!/usr/bin/env python3
from collections import deque, defaultdict
from dataclasses import dataclass
from typing import Generator, Optional, Deque, List, Tuple, DefaultDict
from unittest import TestCase, main


@dataclass
class BinaryTree:
    value: int
    left: Optional['BinaryTree'] = None
    right: Optional['BinaryTree'] = None


def breadth_first_traversal(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    queue = deque([root])  # type: Deque[Optional[BinaryTree]]
    while queue:
        node = queue.popleft()
        if node:
            yield node
            queue.append(node.left)
            queue.append(node.right)


def depth_first_traversal_pre_order(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    yield root
    yield from depth_first_traversal_pre_order(root.left)
    yield from depth_first_traversal_pre_order(root.right)


def depth_first_traversal_in_order(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    yield from depth_first_traversal_in_order(root.left)
    yield root
    yield from depth_first_traversal_in_order(root.right)


def depth_first_traversal_post_order(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    yield from depth_first_traversal_post_order(root.left)
    yield from depth_first_traversal_post_order(root.right)
    yield root


def depth_first_traversal_pre_order_iterative(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    stack = [root]  # type: List[Optional[BinaryTree]]
    while stack:
        node = stack.pop()
        if node:
            yield node
            stack.append(node.right)  # Right before left so we traverse the left subtree first.
            stack.append(node.left)


def depth_first_traversal_in_order_iterative(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    node = root  # type: Optional[BinaryTree]
    stack = []  # type: List[BinaryTree]
    while node or stack:
        if node is not None:
            stack.append(node)
            node = node.left
        else:
            node = stack.pop()
            yield node
            node = node.right


def depth_first_traversal_post_order_iterative(root: Optional[BinaryTree]) -> Generator[BinaryTree, None, None]:
    if root is None:
        return
    unseen = [root]  # type: List[Optional[BinaryTree]]
    seen = []  # type: List[Optional[BinaryTree]]
    while unseen:
        node = unseen.pop()
        seen.append(node)
        if node:
            unseen.append(node.left)  # Left before right so we traverse the right subtree first.
            unseen.append(node.right)
    while seen:
        yield seen.pop()


def leaf_node_traversal(root: Optional[BinaryTree]) -> Generator[Tuple[BinaryTree, int], None, None]:
    if root is None:
        return
    queue = deque([(root, 0)])  # type: Deque[Tuple[BinaryTree, int]]
    while queue:
        node, depth = queue.popleft()
        if node.left or node.right:
            if node.left:
                queue.append((node.left, depth + 1))
            if node.right:
                queue.append((node.right, depth + 1))
        else:
            yield node, depth


def is_full_binary_tree(root: Optional[BinaryTree]) -> bool:
    if root is None:
        return True
    leaf_nodes_by_depth = defaultdict(int)  # type: DefaultDict[int, int]
    for _, depth in leaf_node_traversal(root):
        leaf_nodes_by_depth[depth] += 1
    for leaf_nodes in leaf_nodes_by_depth.values():
        if leaf_nodes % 2 != 0:
            return False
    return True


def is_complete_binary_tree(root: Optional[BinaryTree]) -> bool:
    if root is None:
        return True
    max_depth = None  # type: Optional[int]
    last_depth = None  # type: Optional[int]
    for _, depth in leaf_node_traversal(root):
        if max_depth is None:
            max_depth, last_depth = depth, depth
            continue
        if depth > last_depth or (max_depth - depth) > 1:
            return False
        last_depth = depth
    return True


def is_balanced_binary_tree(root: Optional[BinaryTree]) -> bool:
    if root is None:
        return True
    leaf_node_depths = set(depth for _, depth in leaf_node_traversal(root))
    return len(leaf_node_depths) == 1 or (max(leaf_node_depths) - min(leaf_node_depths)) == 1


def is_perfect_binary_tree(root: Optional[BinaryTree]) -> bool:
    if root is None:
        return True
    leaf_nodes_and_depths = list(leaf_node_traversal(root))
    depths = set(depth for _, depth in leaf_nodes_and_depths)
    if len(depths) != 1:
        return False
    height = list(depths)[0]
    return len(leaf_nodes_and_depths) == 2 ** height


def is_binary_search_tree(root: Optional[BinaryTree]) -> bool:
    last_node = None  # type: Optional[BinaryTree]
    for node in depth_first_traversal_in_order_iterative(root):
        if last_node and last_node.value >= node.value:
            return False
        last_node = node
    return True


class UnitTests(TestCase):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        #         [4]
        #        /  \
        #     [2]   [6]
        #    /  \   / \
        # [1] [3] [5] [7]
        self.tree = BinaryTree(
            4,
            BinaryTree(
                2,
                BinaryTree(
                    1
                ),
                BinaryTree(
                    3
                )
            ),
            BinaryTree(
                6,
                BinaryTree(
                    5
                ),
                BinaryTree(
                    7
                )
            )
        )

    def test_breadth_first_traversal(self) -> None:
        self.assertEqual([4, 2, 6, 1, 3, 5, 7], list(n.value for n in breadth_first_traversal(self.tree) if n is not None))

    def test_depth_first_traversal_pre_order(self) -> None:
        self.assertEqual([4, 2, 1, 3, 6, 5, 7], list(n.value for n in depth_first_traversal_pre_order(self.tree) if n is not None))

    def test_depth_first_traversal_in_order(self) -> None:
        self.assertEqual([1, 2, 3, 4, 5, 6, 7], list(n.value for n in depth_first_traversal_in_order(self.tree) if n is not None))

    def test_depth_first_traversal_post_order(self) -> None:
        self.assertEqual([1, 3, 2, 5, 7, 6, 4], list(n.value for n in depth_first_traversal_post_order(self.tree) if n is not None))

    def test_depth_first_traversal_pre_order_iterative(self) -> None:
        self.assertEqual([4, 2, 1, 3, 6, 5, 7], list(n.value for n in depth_first_traversal_pre_order_iterative(self.tree) if n is not None))

    def test_depth_first_traversal_in_order_iterative(self) -> None:
        self.assertEqual([1, 2, 3, 4, 5, 6, 7], list(n.value for n in depth_first_traversal_in_order_iterative(self.tree) if n is not None))

    def test_depth_first_traversal_post_order_iterative(self) -> None:
        self.assertEqual([1, 3, 2, 5, 7, 6, 4], list(n.value for n in depth_first_traversal_post_order_iterative(self.tree) if n is not None))

    def test_leaf_node_traversal(self) -> None:
        self.assertEqual([(1, 2), (3, 2), (5, 2), (7, 2)], list((node.value, depth) for node, depth in leaf_node_traversal(self.tree)))

    def test_is_full_binary_tree(self) -> None:
        self.assertTrue(is_full_binary_tree(self.tree))
        self.assertFalse(is_full_binary_tree(BinaryTree(
            1,
            BinaryTree(
                2
            )
        )))

    def test_is_complete_binary_tree(self) -> None:
        self.assertTrue(is_complete_binary_tree(self.tree))
        self.assertFalse(is_complete_binary_tree(BinaryTree(
            1,
            BinaryTree(
                2
            ),
            BinaryTree(
                3,
                BinaryTree(
                    4
                ),
                BinaryTree(
                    5
                )
            )
        )))

    def test_is_balanced_binary_tree(self) -> None:
        self.assertTrue(is_balanced_binary_tree(self.tree))
        self.assertFalse(is_balanced_binary_tree(BinaryTree(
            1,
            BinaryTree(
                2,
                BinaryTree(
                    3,
                    BinaryTree(
                        4
                    )
                )
            ),
            BinaryTree(
                5
            )
        )))

    def test_is_perfect_binary_tree(self) -> None:
        self.assertTrue(is_perfect_binary_tree(self.tree))
        self.assertFalse(is_perfect_binary_tree(BinaryTree(
            1,
            BinaryTree(
                2
            )
        )))

    def test_is_binary_search_tree(self) -> None:
        self.assertTrue(is_binary_search_tree(self.tree))
        self.assertFalse(is_binary_search_tree(BinaryTree(
            4,
            BinaryTree(
                2,
                BinaryTree(
                    1,
                ),
                BinaryTree(
                    5,  # Violates BST property (5 > 4).
                ),
            ),
            BinaryTree(
                6,
                BinaryTree(
                    5,
                ),
                BinaryTree(
                    7,
                )
            ),
        )))


if __name__ == '__main__':
    main()