Balanced Trees: AVL & Red-Black
in reviewThis lesson is complete and being validated before release - details may still change.
in reviewThis lesson is complete and being validated before release - details may still change.
avl-tree.ts
import type { Tracer } from "../trace";/** Whether the tree repairs its balance after an insert, or stays a plain binary search tree. */export const AVL_MODES = ["bst", "avl"] as const;export type AvlMode = (typeof AVL_MODES)[number];export const AVL_MIN_KEY = 1;export const AVL_MAX_KEY = 99;/** Keys one tree holds at most, so a stage stays readable. */export const AVL_MAX_SIZE = 32;/** Which side of a node a key goes to, or that it became the root. */export type AvlSide = "left" | "right" | "root";export interface AvlNode { /** Numbered in creation order; a rotation moves nodes but never renumbers them. */ id: number; key: number; left: number | null; right: number | null; /** Nodes on the longest path down from here, this one included; a leaf is 1. */ height: number;}export interface AvlOptions { /** Default avl. */ mode?: AvlMode;}/** * A binary search tree that stays shallow. Every key in a node's left * subtree is smaller than the node's key and every key in its right * subtree larger, so a lookup descends one comparison per level and * costs the height of the tree. A plain tree takes keys as they come, * and sorted input hangs them in a chain: height n, lookups O(n). An * AVL tree keeps, at every node, the heights of the two subtrees * within one of each other. An insert descends, attaches a leaf, and * walks back up refreshing heights; the first node whose balance * reaches two is rotated: the taller child becomes the parent and the * node drops to its other side, with one rotation when the weight is * on the outside and two when it is on the inside. The imbalance is * repaired at that node and nothing above it needs another rotation, * so an insert does O(1) rotations, and the invariant bounds the height * by about 1.44 log2 n. * * Time: O(log n) per insert or search in AVL mode, O(height) up to * O(n) as a plain tree. Space: O(n). */export class AvlTree { readonly mode: AvlMode; private readonly nodes = new Map<number, AvlNode>(); private nextId = 0; private rotations = 0; root: number | null = null; constructor(options: AvlOptions = {}, private readonly tracer?: Tracer) { this.mode = options.mode ?? "avl"; if (!AVL_MODES.includes(this.mode)) throw new RangeError(`mode must be one of ${AVL_MODES.join(", ")}`); } /** How many keys the tree holds. */ get size(): number { return this.nodes.size; } /** Rotations performed so far. */ get rotationCount(): number { return this.rotations; } /** Every node, by id; copies. */ allNodes(): readonly AvlNode[] { return [...this.nodes.values()].map((node) => ({ ...node })); } /** Nodes on the longest root-to-leaf path; 0 when empty. */ height(): number { return this.root === null ? 0 : this.node(this.root).height; } /** The keys in ascending order. */ keys(): number[] { const out: number[] = []; const walk = (id: number | null): void => { if (id === null) return; const node = this.node(id); walk(node.left); out.push(node.key); walk(node.right); }; walk(this.root); return out; } /** Adds a key the tree does not hold yet: descend, attach a leaf, then refresh and (in AVL mode) rebalance on the way back up. */ insert(key: number): void { this.checkKey(key); if (this.nodes.size >= AVL_MAX_SIZE) throw new RangeError(`the tree holds at most ${AVL_MAX_SIZE} keys`); if (this.root === null) { this.root = this.attach(key, null, "root"); return; } this.root = this.insertAt(this.root, key); } /** Whether the tree holds the key, reading one node per level. */ search(key: number): boolean { this.checkKey(key); let reads = 0; let at = this.root; while (at !== null) { const node = this.node(at); reads++; if (key === node.key) { // ▸ the node read holds the key: the lookup cost as many reads as the node's depth this.tracer?.emit("found", { key, node: at, reads }); return true; } const dir = this.descend(at, node.key, key); at = dir === "left" ? node.left : node.right; } // ▸ the side the key would be on is empty: the tree does not hold it this.tracer?.emit("missing", { key, reads }); return false; } /** Which way a key goes past a node, the one step every search and insert repeats per level. */ private descend(node: number, at: number, key: number): AvlSide { const dir: AvlSide = key < at ? "left" : "right"; // ▸ compare the key with this node's: smaller goes left, larger goes right, one comparison per level this.tracer?.emit("descend", { node, at, key, dir }); return dir; } /** The subtree under `id` after the key goes in, which is a different node when a rotation lifted a child. */ private insertAt(id: number, key: number): number { const node = this.node(id); if (key === node.key) throw new RangeError(`the tree already holds ${key}`); const dir = this.descend(id, node.key, key); const child = dir === "left" ? node.left : node.right; const below = child === null ? this.attach(key, id, dir) : this.insertAt(child, key); if (dir === "left") node.left = below; else node.right = below; const balance = this.refresh(node); const leftMinusRightHeight = this.refresh(node); // ▸ back on the way up, this node's height is one more than its taller child; its balance is left height minus right this.tracer?.emit("update", { node: id, key: node.key, height: node.height, balance }); this.tracer?.emit("update", { node: id, key: node.key, height: node.height, leftMinusRightHeight }); return this.mode === "avl" ? this.rebalance(id, balance) : id; return this.mode === "avl" ? this.rebalance(id, leftMinusRightHeight) : id; } /** A new leaf for the key, hung on `parent`'s `dir` side (the root when parent is null). */ private attach(key: number, parent: number | null, dir: AvlSide): number { const id = this.nextId++; this.nodes.set(id, { id, key, left: null, right: null, height: 1 }); // ▸ the side the key wanted is empty: hang it there as a new leaf this.tracer?.emit("attach", { node: id, key, parent, dir }); return id; } /** Restores the invariant at a node whose balance has reached two: one rotation for an outside case, two for an inside one. */ private rebalance(id: number, balance: number): number { private rebalance(id: number, leftMinusRightHeight: number): number { const node = this.node(id); if (balance > 1) { if (leftMinusRightHeight > 1) { // Left-heavy. Inside case (left-right): the left child leans right, so straighten it first. if (this.balanceOf(node.left!) < 0) node.left = this.rotateLeft(node.left!); return this.rotateRight(id); } if (balance < -1) { if (leftMinusRightHeight < -1) { // Right-heavy. Inside case (right-left): the right child leans left, so straighten it first. if (this.balanceOf(node.right!) > 0) node.right = this.rotateRight(node.right!); return this.rotateLeft(id); } return id; } /** The right child rises and the node drops to its left; the pivot's old left subtree moves across. */ private rotateLeft(id: number): number { const node = this.node(id); const pivot = this.node(node.right!); const risingChild = this.node(node.right!); node.right = pivot.left; node.right = risingChild.left; pivot.left = id; risingChild.left = id; this.refresh(node); this.refresh(pivot); this.refresh(risingChild); this.rotations++; // ▸ the right child rises to take this node's place and the node drops to its left; the child's old left subtree crosses over this.tracer?.emit("rotate-left", { node: id, nodeKey: node.key, pivot: pivot.id, pivotKey: pivot.key }); this.tracer?.emit("rotate-left", { node: id, nodeKey: node.key, risingChild: risingChild.id, pivotKey: risingChild.key }); return pivot.id; return risingChild.id; } /** The left child rises and the node drops to its right; the pivot's old right subtree moves across. */ private rotateRight(id: number): number { const node = this.node(id); const pivot = this.node(node.left!); const risingChild = this.node(node.left!); node.left = pivot.right; node.left = risingChild.right; pivot.right = id; risingChild.right = id; this.refresh(node); this.refresh(pivot); this.refresh(risingChild); this.rotations++; // ▸ the left child rises to take this node's place and the node drops to its right; the child's old right subtree crosses over this.tracer?.emit("rotate-right", { node: id, nodeKey: node.key, pivot: pivot.id, pivotKey: pivot.key }); this.tracer?.emit("rotate-right", { node: id, nodeKey: node.key, risingChild: risingChild.id, pivotKey: risingChild.key }); return pivot.id; return risingChild.id; } /** Recomputes a node's height from its children and returns its balance: left height minus right height. */ private refresh(node: AvlNode): number { const left = this.heightOf(node.left); const right = this.heightOf(node.right); node.height = 1 + Math.max(left, right); return left - right; } private balanceOf(id: number): number { const node = this.node(id); return this.heightOf(node.left) - this.heightOf(node.right); } private heightOf(id: number | null): number { return id === null ? 0 : this.node(id).height; } private node(id: number): AvlNode { const node = this.nodes.get(id); if (!node) throw new RangeError(`no node ${id}`); return node; } private checkKey(key: number): void { if (!Number.isInteger(key) || key < AVL_MIN_KEY || key > AVL_MAX_KEY) throw new RangeError(`key must be an integer in [${AVL_MIN_KEY}, ${AVL_MAX_KEY}]`); }}10 keys in a AVL tree 4 levels high, where 4 is the fewest any tree of that many keys could have; 6 rotations along the way. The lookup of 27 found it in 3 reads.