S028 · Trees & BSTs — Traversal (BFS/DFS)

"""trees.py — build a BST, traverse 4 ways, validate, and find k-th smallest.
 
Run:  python trees.py
"""
from __future__ import annotations
from collections import deque
from typing import Optional
 
 
class Node:
    __slots__ = ("val", "left", "right")
 
    def __init__(self, val: int, left: "Node | None" = None, right: "Node | None" = None):
        self.val = val
        self.left = left
        self.right = right
 
 
# ------------------------------------------------------------------
# 1) BST insert + search
# ------------------------------------------------------------------
def bst_insert(root: Node | None, val: int) -> Node:
    if root is None:
        return Node(val)
    if val < root.val:
        root.left = bst_insert(root.left, val)
    elif val > root.val:
        root.right = bst_insert(root.right, val)
    # duplicates: silently ignored
    return root
 
 
def bst_search(root: Node | None, val: int) -> bool:
    node = root
    while node:
        if val == node.val:
            return True
        node = node.left if val < node.val else node.right
    return False
 
 
# ------------------------------------------------------------------
# 2) The four traversals
# ------------------------------------------------------------------
def pre_order(root: Node | None, out: list[int] | None = None) -> list[int]:
    out = out if out is not None else []
    if root:
        out.append(root.val)
        pre_order(root.left, out)
        pre_order(root.right, out)
    return out
 
 
def in_order(root: Node | None, out: list[int] | None = None) -> list[int]:
    out = out if out is not None else []
    if root:
        in_order(root.left, out)
        out.append(root.val)
        in_order(root.right, out)
    return out
 
 
def post_order(root: Node | None, out: list[int] | None = None) -> list[int]:
    out = out if out is not None else []
    if root:
        post_order(root.left, out)
        post_order(root.right, out)
        out.append(root.val)
    return out
 
 
def level_order(root: Node | None) -> list[list[int]]:
    """BFS, grouping values by level."""
    if not root:
        return []
    levels: list[list[int]] = []
    q: deque[Node] = deque([root])
    while q:
        level = []
        for _ in range(len(q)):                      # snapshot the current level size
            node = q.popleft()
            level.append(node.val)
            if node.left:
                q.append(node.left)
            if node.right:
                q.append(node.right)
        levels.append(level)
    return levels
 
 
# ------------------------------------------------------------------
# 3) Tree properties — height, size, is-BST
# ------------------------------------------------------------------
def height(root: Node | None) -> int:
    if root is None:
        return 0
    return 1 + max(height(root.left), height(root.right))
 
 
def size(root: Node | None) -> int:
    if root is None:
        return 0
    return 1 + size(root.left) + size(root.right)
 
 
def is_valid_bst(root: Node | None, low: float = float("-inf"), high: float = float("inf")) -> bool:
    """Every node must lie strictly in (low, high). Recurse tightening the bounds."""
    if root is None:
        return True
    if not (low < root.val < high):
        return False
    return is_valid_bst(root.left, low, root.val) and is_valid_bst(root.right, root.val, high)
 
 
# ------------------------------------------------------------------
# 4) k-th smallest — leverage in-order + early stop
# ------------------------------------------------------------------
def kth_smallest(root: Node | None, k: int) -> Optional[int]:
    """Return the k-th smallest value in the BST, or None. O(h + k) time."""
    stack: list[Node] = []
    node = root
    count = 0
    while node or stack:
        while node:
            stack.append(node)
            node = node.left
        node = stack.pop()
        count += 1
        if count == k:
            return node.val
        node = node.right
    return None
 
 
# ------------------------------------------------------------------
# 5) Compare recursive vs iterative BFS on a balanced random tree
# ------------------------------------------------------------------
def build_bst_from(values: list[int]) -> Node | None:
    root: Node | None = None
    for v in values:
        root = bst_insert(root, v)
    return root
 
 
if __name__ == "__main__":
    # Build the tree:
    #         50
    #        /  \
    #      30    70
    #     / \   / \
    #   20  40 60 80
    root = build_bst_from([50, 30, 70, 20, 40, 60, 80])
 
    print("pre_order :", pre_order(root))
    print("in_order  :", in_order(root))            # sorted for a valid BST
    print("post_order:", post_order(root))
    print("level_order:", level_order(root))
    print("height    :", height(root))              # 3
    print("size      :", size(root))                # 7
    print("is_valid_bst :", is_valid_bst(root))     # True
    print("kth_smallest k=3 :", kth_smallest(root, 3))   # 40
    print("bst_search 60 :", bst_search(root, 60))
    print("bst_search 65 :", bst_search(root, 65))
 
    # Demonstrate the balance problem — insert sorted data
    bad = build_bst_from([1, 2, 3, 4, 5, 6, 7])
    print(f"height of BAD tree from sorted input : {height(bad)}  (should be ~3, actually 7 — a linked list)")