Binary search trees: search, insert, delete

"""A binary search tree: search, insert and the three cases of delete, with the tree printed after each step."""


class Node:
    def __init__(self, key):
        self.key, self.left, self.right = key, None, None


def search(root, key, path):
    """Walks down from the root; path collects the keys compared on the way."""
    node = root
    while node is not None:
        path.append(node.key)
        if key == node.key:
            return node
        node = node.left if key < node.key else node.right
    return None


def insert(node, key):
    """Returns the subtree with key added; an existing key is left alone (each key is stored once)."""
    if node is None:
        return Node(key)
    if key < node.key:
        node.left = insert(node.left, key)
    elif key > node.key:
        node.right = insert(node.right, key)
    return node


def delete(node, key):
    """Returns the subtree with key removed."""
    if node is None:
        return None                              # the key is not in the tree
    if key < node.key:
        node.left = delete(node.left, key)
    elif key > node.key:
        node.right = delete(node.right, key)
    elif node.left is None:                      # no child, or only a right child: the child takes the node's place
        return node.right
    elif node.right is None:                     # only a left child
        return node.left
    else:                                        # two children: take the inorder successor's key, delete it below
        successor = node.right
        while successor.left is not None:
            successor = successor.left
        node.key = successor.key
        node.right = delete(node.right, successor.key)
    return node


def inorder(node, out):
    if node:
        inorder(node.left, out)
        out.append(node.key)
        inorder(node.right, out)
    return out


def shape(node):
    """The tree as nested brackets: key(left, right), with - for an empty child."""
    if node is None:
        return "-"
    if node.left is None and node.right is None:
        return str(node.key)
    return f"{node.key}({shape(node.left)}, {shape(node.right)})"


root = None
for k in (50, 30, 70, 20, 40, 60, 80, 55, 65):
    root = insert(root, k)
print("tree:", shape(root))
for key in (65, 58):
    path = []
    found = search(root, key, path)
    print(f"search {key}: compared {path}, {'found' if found else 'not found'}")
root = insert(root, 58)
print("after insert 58:", shape(root))
for key, case in ((20, "a leaf"), (55, "one child"), (50, "two children")):
    root = delete(root, key)
    print(f"after delete {key} ({case}):", shape(root))
print("inorder:", inorder(root, []))

tree: 50(30(20, 40), 70(60(55, 65), 80))
search 65: compared [50, 70, 60, 65], found
search 58: compared [50, 70, 60, 55], not found
after insert 58: 50(30(20, 40), 70(60(55(-, 58), 65), 80))
after delete 20 (a leaf): 50(30(-, 40), 70(60(55(-, 58), 65), 80))
after delete 55 (one child): 50(30(-, 40), 70(60(58, 65), 80))
after delete 50 (two children): 58(30(-, 40), 70(60(-, 65), 80))
inorder: [30, 40, 58, 60, 65, 70, 80]

"""The height of a binary search tree after 4,096 insertions: keys in random order against keys in sorted order."""
import random


class Node:
    __slots__ = ("key", "left", "right")

    def __init__(self, key):
        self.key, self.left, self.right = key, None, None


def insert(root, key):
    """Walks down from the root without recursion and attaches the key where the search fell off the tree."""
    if root is None:
        return Node(key)
    node = root
    while True:
        if key < node.key:
            if node.left is None:
                node.left = Node(key)
                return root
            node = node.left
        elif key > node.key:
            if node.right is None:
                node.right = Node(key)
                return root
            node = node.right
        else:
            return root                      # already present: this tree keeps each key once


def height_and_depths(root):
    """Height in edges and the average depth of a node (the comparisons a search for a stored key makes, minus one)."""
    total, height, level, depth = 0, -1, [root], 0
    while level:
        height += 1
        total += depth * len(level)
        level = [c for n in level for c in (n.left, n.right) if c]
        depth += 1
    return height, total


N = 4_096
rng = random.Random(4096)
keys = list(range(N))
shuffled = keys[:]
rng.shuffle(shuffled)
for label, order in (("random order", shuffled), ("sorted order", keys)):
    root = None
    for k in order:
        root = insert(root, k)
    h, total = height_and_depths(root)
    print(f"{N:,} keys in {label}: height {h:,}, average depth {total / N:,.2f}")

heights, depths = [], []
for _ in range(200):
    rng.shuffle(shuffled)
    root = None
    for k in shuffled:
        root = insert(root, k)
    h, total = height_and_depths(root)
    heights.append(h)
    depths.append(total / N)
print(f"200 random orders: height from {min(heights)} to {max(heights)}, mean {sum(heights) / 200:.2f};"
      f" average depth {sum(depths) / 200:.2f}")
print(f"for comparison: floor(log2 {N:,}) = {N.bit_length() - 1}, the lowest possible height; {N - 1:,}, the highest")

4,096 keys in random order: height 23, average depth 13.09
4,096 keys in sorted order: height 4,095, average depth 2,047.50
200 random orders: height from 23 to 36, mean 26.61; average depth 13.82
for comparison: floor(log2 4,096) = 12, the lowest possible height; 4,095, the highest