Red-black trees and B-trees (concepts)

"""A red-black tree: insertion with each repair printed, a check of the rules, and the 2-3-4 tree it stands for."""

RED, BLACK = "red", "black"


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

    def __init__(self, key, color, nil):
        self.key, self.color = key, color
        self.left = self.right = self.parent = nil


class RedBlackTree:
    def __init__(self):
        self.nil = Node(None, BLACK, None)  # one shared black leaf stands for every empty subtree
        self.nil.left = self.nil.right = self.nil.parent = self.nil
        self.root = self.nil
        self.rotations = 0
        self.log = []

    def rotate_left(self, x):
        y = x.right
        x.right = y.left
        if y.left is not self.nil:
            y.left.parent = x
        self._replace(x, y)
        y.left, x.parent = x, y
        self.rotations += 1

    def rotate_right(self, y):
        x = y.left
        y.left = x.right
        if x.right is not self.nil:
            x.right.parent = y
        self._replace(y, x)
        x.right, y.parent = y, x
        self.rotations += 1

    def _replace(self, old, new):
        """Puts new where old hangs from its parent."""
        new.parent = old.parent
        if old.parent is self.nil:
            self.root = new
        elif old is old.parent.left:
            old.parent.left = new
        else:
            old.parent.right = new

    def insert(self, key):
        parent, node = self.nil, self.root
        while node is not self.nil:
            if key == node.key:
                return  # keys are unique
            parent, node = node, (node.left if key < node.key else node.right)
        z = Node(key, RED, self.nil)  # a new node is red, so no path gains a black node
        z.parent = parent
        if parent is self.nil:
            self.root = z
        elif key < parent.key:
            parent.left = z
        else:
            parent.right = z
        self._repair(z)

    def _repair(self, z):
        while z.parent.color == RED:  # two reds in a row: the only rule an insertion can break
            p, g = z.parent, z.parent.parent
            uncle = g.right if p is g.left else g.left
            if uncle.color == RED:  # recolour and move the problem two levels up
                p.color = uncle.color = BLACK
                g.color = RED
                self.log.append(f"red uncle at {g.key}: recolour")
                z = g
                continue
            if p is g.left:
                if z is p.right:  # zig-zag: rotate it into a straight line first
                    self.rotate_left(p)
                    z, p = p, z
                    self.log.append(f"black uncle, zig-zag at {g.key}: two rotations")
                else:
                    self.log.append(f"black uncle, line at {g.key}: one rotation")
                self.rotate_right(g)
            else:
                if z is p.left:
                    self.rotate_right(p)
                    z, p = p, z
                    self.log.append(f"black uncle, zig-zag at {g.key}: two rotations")
                else:
                    self.log.append(f"black uncle, line at {g.key}: one rotation")
                self.rotate_left(g)
            p.color, g.color = BLACK, RED  # the middle key goes up black, the old grandparent becomes red
        if self.root.color == RED:
            self.root.color = BLACK
            self.log.append("root painted black")

    def check(self):
        """Black nodes on every path from the root to an empty leaf; fails on the first broken rule."""
        assert self.root.color == BLACK, "the root is red"

        def walk(n, low, high):
            if n is self.nil:
                return 0
            assert low < n.key < high, f"{n.key} is out of order"
            if n.color == RED:
                assert n.left.color == BLACK and n.right.color == BLACK, f"red {n.key} has a red child"
            left, right = walk(n.left, low, n.key), walk(n.right, n.key, high)
            assert left == right, f"paths below {n.key} hold {left} and {right} black nodes"
            return left + (n.color == BLACK)

        return walk(self.root, float("-inf"), float("inf"))

    def height(self, n=None):
        n = self.root if n is None else n
        return -1 if n is self.nil else 1 + max(self.height(n.left), self.height(n.right))

    def sideways(self, n=None, depth=0, out=None):
        n, out = (self.root if n is None else n), ([] if out is None else out)
        if n is not self.nil:
            self.sideways(n.right, depth + 1, out)
            out.append("    " * depth + f"{n.key} {n.color}")
            self.sideways(n.left, depth + 1, out)
        return out

    def as_234(self):
        """The 2-3-4 tree level by level: each black node merged with its red children into one node."""
        levels, level = [], [self.root] if self.root is not self.nil else []
        while level:
            row, below = [], []
            for b in level:
                keys = []
                for c in (b.left, b, b.right):
                    if c is b:
                        keys.append(b.key)
                    elif c.color == RED:
                        keys.append(c.key)
                        below += [c.left, c.right]
                    else:
                        below.append(c)
                row.append("[" + " ".join(map(str, keys)) + "]")
            levels.append(" ".join(row))
            level = [x for x in below if x is not self.nil]
        return levels


if __name__ == "__main__":
    t = RedBlackTree()
    for key in (10, 20, 30, 15, 25, 5, 1, 12, 14, 13):
        t.log.clear()
        t.insert(key)
        t.check()  # every rule holds after every insertion
        print(f"insert {key:2}: {'; '.join(t.log) or 'parent is black: nothing to do'}")
    print(f"10 keys: height {t.height()}, {t.check()} black nodes on every root-to-leaf path, {t.rotations} rotations")
    print("\n".join(t.sideways()))
    print("the same keys as a 2-3-4 tree, level by level:")
    for row in t.as_234():
        print("   ", row)

insert 10: root painted black
insert 20: parent is black: nothing to do
insert 30: black uncle, line at 10: one rotation
insert 15: red uncle at 20: recolour; root painted black
insert 25: parent is black: nothing to do
insert  5: parent is black: nothing to do
insert  1: red uncle at 10: recolour
insert 12: parent is black: nothing to do
insert 14: black uncle, zig-zag at 15: two rotations
insert 13: red uncle at 14: recolour; black uncle, zig-zag at 20: two rotations
10 keys: height 3, 2 black nodes on every root-to-leaf path, 5 rotations
        30 black
            25 red
    20 red
        15 black
14 black
            13 red
        12 black
    10 red
        5 black
            1 red
the same keys as a 2-3-4 tree, level by level:
    [10 14 20]
    [1 5] [12 13] [15] [25 30]

"""A B-tree of order 4 (at most 3 keys and 4 children per node) built from 20 keys, with every split printed."""

import random


class BNode:
    def __init__(self, keys=None, children=None):
        self.keys = keys or []
        self.children = children or []  # empty for a leaf


def insert(node, key, log, order):
    """Inserts key below node; returns (median, new right sibling) when node splits, else None."""
    i = sum(k < key for k in node.keys)  # the child (or slot) to the right of every smaller key
    if i < len(node.keys) and node.keys[i] == key:
        return None  # keys are unique
    if not node.children:
        node.keys.insert(i, key)
    else:
        split = insert(node.children[i], key, log, order)
        if split:
            median, right = split
            node.keys.insert(i, median)
            node.children.insert(i + 1, right)
    if len(node.keys) < order:
        return None
    mid = len(node.keys) // 2  # an overfull node holds `order` keys: keep the left half, send the middle key up
    median = node.keys[mid]
    right = BNode(node.keys[mid + 1 :], node.children[mid + 1 :])
    log.append(f"[{' '.join(map(str, node.keys))}] splits, {median} moves up")
    node.keys, node.children = node.keys[:mid], node.children[: mid + 1]
    return median, right


def add(root, key, log, order=4):
    split = insert(root, key, log, order)
    if split:  # the root itself split: a new root makes the tree one level taller
        median, right = split
        root = BNode([median], [root, right])
        log.append(f"new root [{median}]")
    return root


def levels(root):
    rows, level = [], [root]
    while level:
        rows.append(" ".join("[" + " ".join(map(str, n.keys)) + "]" for n in level))
        level = [c for n in level for c in n.children]
    return rows


def build(keys, order=4):
    root = BNode()
    for key in keys:
        root = add(root, key, [], order)
    return root


if __name__ == "__main__":
    keys = list(range(1, 21))
    random.Random(4).shuffle(keys)
    root = BNode()
    for key in keys:
        log = []
        root = add(root, key, log)
        if log:
            print(f"insert {key:2}: {'; '.join(log)}")
    print(f"{len(keys)} keys inserted in the order {keys}")
    for depth, row in enumerate(levels(root)):
        print(f"level {depth}: {row}")

insert 15: [11 15 18 20] splits, 18 moves up; new root [18]
insert 17: [6 11 15 17] splits, 15 moves up
insert  5: [5 6 11 12] splits, 11 moves up
insert  7: [5 6 7 9] splits, 7 moves up; [7 11 15 18] splits, 15 moves up; new root [15]
insert  2: [1 2 5 6] splits, 5 moves up
insert  4: [1 2 3 4] splits, 3 moves up; [3 5 7 11] splits, 7 moves up
20 keys inserted in the order [18, 20, 11, 15, 6, 19, 17, 12, 5, 9, 7, 1, 14, 2, 3, 16, 13, 4, 10, 8]
level 0: [7 15]
level 1: [3 5] [11] [18]
level 2: [1 2] [4] [6] [8 9 10] [12 13 14] [16 17] [19 20]

"""How many nodes a lookup visits: a B-tree with 4 KB pages against binary trees, for 10^3 to 10^9 keys."""

import math

PAGE, ENTRY = 4096, 16  # bytes per page; one entry = an 8-byte key and an 8-byte page number
fanout = PAGE // ENTRY  # children per full node: 256


def btree_levels(n, children):
    """Levels of a B-tree whose nodes are all full (the best case) and half full (the worst a B-tree allows)."""
    return math.ceil(math.log(n + 1, children))


print(f"page {PAGE} bytes, entry {ENTRY} bytes: up to {fanout} children per node, at least {fanout // 2}")
print("nodes one lookup visits, at most:")
print(f"{'keys':>15} {'perfect binary':>15} {'red-black worst':>16} {'B-tree, full':>13} {'B-tree, half full':>18}")
for n in (10**3, 10**6, 10**9):
    perfect = math.ceil(math.log2(n + 1))
    rb_worst = math.floor(2 * math.log2(n + 1))
    print(f"{n:>15,} {perfect:>15} {rb_worst:>16} {btree_levels(n, fanout):>13} {btree_levels(n, fanout // 2):>18}")

page 4096 bytes, entry 16 bytes: up to 256 children per node, at least 128
nodes one lookup visits, at most:
           keys  perfect binary  red-black worst  B-tree, full  B-tree, half full
          1,000              10               19             2                  2
      1,000,000              20               39             3                  3
  1,000,000,000              30               59             4                  5