Skip to content
Browse chapters

Chapter 17 — Augmenting Data Structures

CLRS, fourth edition · Lean 4 formalization

The proofs below use the models and assumptions described in the scope and implementation notes.

Imports

17.1. Dynamic Order Statistics

This section completes the fourth-edition §17.1 boundary for dynamic order-statistic trees. The legacy CLRSLean.Chapter_14 development already provides OS-SELECT (OSRBTree.osSelect?) and the subtree-size augmentation through executable red-black insertion and deletion. The missing operation is OS-RANK — given a key, return its rank (the number of stored keys strictly smaller than it).

We add OSRBTree.osRank (using cached sizes) and the ideal OSRBTree.rankOf (using recomputed sizes), prove they agree on well-sized trees. The RankCardinality companion proves that under strict BST ordering, both ranks equal the cardinality of keys strictly below the query, including when that query is absent. We also prove OS-RANK runs in O(log n) pointer operations via the height of the underlying red-black tree.

Main results:

  • Definition OSRBTree.osRank / OSRBTree.rankOf: OS-RANK with cached and recomputed sizes.

  • Theorem OSRBTree.osRank_eq_rankOf_of_wellSized: cached and ideal ranks agree on well-sized trees.

  • Theorem OSRBTree.osRankCost_le_height: the pointer cost of OS-RANK is bounded by the tree height.

  • Theorem OSRBTree.osRankCost_log_bound: OS-RANK runs in O(log n) on a red-black-shaped tree.

namespace CLRSnamespace Chapter14namespace OSRBTreeopen CLRS.Chapter13 (RBTree)

OS-RANK

The ideal rank of x: the number of stored keys strictly smaller than x, computed with recomputed (mathematical) subtree sizes.

def rankOf : OSRBTree → Nat → Nat | empty, _ => 0 | node _ l k _ r, x => if x < k then rankOf l x else if x > k then realSize l + 1 + rankOf r x else realSize l

OS-RANK using the cached size fields.

def osRank : OSRBTree → Nat → Nat | empty, _ => 0 | node _ l k _ r, x => if x < k then osRank l x else if x > k then storedSize l + 1 + osRank r x else storedSize l

The pointer-operation cost of OS-RANK: one node read per level of the descent (a root-to-leaf path).

def osRankCost : OSRBTree → Nat → Nat | empty, _ => 0 | node _ l k _ r, x => if x < k then 1 + osRankCost l x else if x > k then 1 + osRankCost r x else 1

The cached OS-RANK agrees with the ideal rank on well-sized trees.

theorem osRank_eq_rankOf_of_wellSized {t : OSRBTree} {x : Nat} (h : WellSized t) : osRank t x = rankOf t x := by induction t generalizing x with | empty => rfl | node c l k s r ihl ihr => obtain ⟨hL, hR, _⟩ := h have hlSize : storedSize l = realSize l := storedSize_eq_realSize_of_wellSized hL by_cases hlt : x < k · simp [osRank, rankOf, hlt, ihl hL] · by_cases hgt : x > k · simp [osRank, rankOf, hlt, hgt, hlSize, ihr hR] · have heq : x = k := by omega simp [osRank, rankOf, This simp argument is unused: hlt Hint: Omit it from the simp argument list. simp [osRank, rankOf, hl̵t̵,̵ ̵h̵gt, heq, hlSize] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`hlt, This simp argument is unused: hgt Hint: Omit it from the simp argument list. simp [osRank, rankOf, hlt, hg̵t̵,̵ ̵h̵eq, hlSize] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`hgt, heq, hlSize]

Erasing the size field preserves the node count: the Chapter 13 size of the erasure equals the mathematical realSize.

theorem toRB_size (t : OSRBTree) : RBTree.size (toRB t) = realSize t := by induction t with | empty => rfl | node c l k s r ihl ihr => simp [toRB, RBTree.size, realSize, ihl, ihr]; omega

The pointer cost of OS-RANK is bounded by the height of the underlying tree.

theorem osRankCost_le_height (t : OSRBTree) (x : Nat) : osRankCost t x ≤ RBTree.height (toRB t) + 1 := by induction t generalizing x with | empty => simp [osRankCost, toRB, RBTree.height] | node c l k s r ihl ihr => simp only [osRankCost, toRB, RBTree.height] by_cases hlt : x < k · simp [hlt] have ih := ihl x have hmax : RBTree.height (toRB l) ≤ max (RBTree.height (toRB l)) (RBTree.height (toRB r)) := Nat.le_max_left _ _ omega · by_cases hgt : x > k · simp [hlt, hgt] have ih := ihr x have hmax : RBTree.height (toRB r) ≤ max (RBTree.height (toRB l)) (RBTree.height (toRB r)) := Nat.le_max_right _ _ omega · simp [hlt, hgt]

OS-RANK runs in O(log n). On a red-black-shaped tree with n nodes, OS-RANK performs at most 2 log₂(n+1) + 1 pointer operations.

theorem osRankCost_log_bound (t : OSRBTree) (x : Nat) (hShape : RBTree.RedBlackShape (toRB t)) : osRankCost t x ≤ 2 * Nat.log 2 (realSize t + 1) + 1 := by have hh := RBTree.height_log_bound (toRB t) hShape have hc := osRankCost_le_height t x rw [toRB_size t] at hh omega
end OSRBTreeend Chapter14end CLRS

Definitions and proofs

CLRSLean.FourthEdition.Chapter_17.Section_17_1_Dynamic_Order_Statistics.RankCardinality

Order-statistic rank as a set cardinality

BST ordering and cached-size correctness are separate requirements. Ordering identifies the recursive rank with the number of keys strictly below the query; well-sizedness then transfers that interpretation to the cached query. The query need not be present, and the answer is a zero-based insertion rank.

namespace CLRS.Chapter14.OSRBTreeopen CLRS.Chapter13 (RBTree)theorem toRB_keys (t : OSRBTree) : RBTree.keys (toRB t) = keys t := by induction t with | empty => rfl | node c l k s r ihl ihr => simp [toRB, keys, RBTree.keys, ihl, ihr]theorem keys_length (t : OSRBTree) : (keys t).length = realSize t := by induction t with | empty => rfl | node c l k s r ihl ihr => simp [keys, realSize, ihl, ihr]; omega theorem keys_nodup_of_bst {t : OSRBTree} (h : RBTree.BST (toRB t)) : (keys t).Nodup := by have hs := (RBTree.bst_iff_sorted (toRB t)).mp h rw [RBTree.sorted, toRB_keys] at hs exact hs.imp (fun hlt => Nat.ne_of_lt hlt)private theorem filter_lt_all (ys : List Nat) (x : Nat) (h : ∀ y ∈ ys, y < x) : ys.filter (fun y => decide (y < x)) = ys := by apply List.filter_eq_self.mpr intro y hy simpa using h y hyprivate theorem filter_lt_none (ys : List Nat) (x : Nat) (h : ∀ y ∈ ys, x ≤ y) : ys.filter (fun y => decide (y < x)) = [] := by apply List.filter_eq_nil_iff.mpr intro y hy simpa using Nat.not_lt.mpr (h y hy)

Recursive rank counts the strict lower-key prefix, even for absent queries.

theorem rankOf_eq_filter_length {t : OSRBTree} (h : RBTree.BST (toRB t)) (x : Nat) : rankOf t x = ((keys t).filter (fun y => decide (y < x))).length := by induction t with | empty => simp [rankOf, keys] | node c l k s r ihl ihr => rcases h with ⟨hl, hr, hleft, hright⟩ have hlkeys : ∀ y ∈ keys l, y < k := by intro y hy exact hleft y ((inTree_toRB y l).mpr hy) have hrkeys : ∀ y ∈ keys r, k < y := by intro y hy exact hright y ((inTree_toRB y r).mpr hy) by_cases hx : x < k · have hn := filter_lt_none (keys r) x (by intro y hy; exact Nat.le_of_lt (lt_trans hx (hrkeys y hy))) simp [rankOf, keys, hx, show ¬k < x by omega, List.filter_append, hn, ihl hl] · by_cases hk : k < x · have ha := filter_lt_all (keys l) x (by intro y hy; exact lt_trans (hlkeys y hy) hk) simp [rankOf, keys, hx, hk, List.filter_append, ha, keys_length, ihr hr] omega · have heq : x = k := by omega subst x have ha := filter_lt_all (keys l) k hlkeys have hn := filter_lt_none (keys r) k (by intro y hy; exact Nat.le_of_lt (hrkeys y hy)) simp [rankOf, keys, List.filter_append, ha, hn, keys_length]

BST trees have distinct keys, so the list-prefix count is a set cardinality.

theorem rankOf_eq_card_lt {t : OSRBTree} (h : RBTree.BST (toRB t)) (x : Nat) : rankOf t x = ((keys t).toFinset.filter (fun y => y < x)).card := by rw [rankOf_eq_filter_length h] have hn := (keys_nodup_of_bst h).filter (fun y => decide (y < x)) rw [← List.toFinset_card_of_nodup hn] congr 1 ext y simp

Semantic cardinal-rank contract for the cached query.

theorem osRank_eq_card_lt {t : OSRBTree} (hs : WellSized t) (hb : RBTree.BST (toRB t)) (x : Nat) : osRank t x = ((keys t).toFinset.filter (fun y => y < x)).card := by rw [osRank_eq_rankOf_of_wellSized hs, rankOf_eq_card_lt hb]
end CLRS.Chapter14.OSRBTree

CLRSLean.Chapter_14.Section_14_1_Order_Statistic_Trees

CLRS Section 14.1 - Order-statistic trees

This section gives the first augmentation proof for CLRS-Lean. An order-statistic tree stores, at every node, a size field intended to equal the number of nodes in that subtree. The operation osSelect? uses the stored size of the left child to implement rank selection.

The file separates the executable augmented operation from its ideal mathematical specification:

  • storedSize reads the cached field;

  • realSize recomputes the mathematical subtree size;

  • WellSized says every cached field is correct;

  • rankSelect? is the ideal selector using recomputed sizes;

  • osSelect? is the augmented selector using cached sizes.

Main results:

  • Theorem storedSize_eq_realSize_of_wellSized: a well-sized tree has a correct root size field.

  • Theorem realSize_recomputeSizes: recomputing cached size fields preserves the mathematical subtree size.

  • Theorem recomputeSizes_wellSized: recomputing size fields establishes the augmentation invariant.

  • Theorem keys_recomputeSizes: recomputing size fields preserves the inorder key sequence.

  • Theorem rankSelect?_recomputeSizes: recomputing size fields preserves the ideal rank-selection result.

  • Theorems keys_rotateLeft and keys_rotateRight: rotations preserve the inorder key sequence.

  • Theorems rotateLeft_wellSized and rotateRight_wellSized: rotations with local size recomputation preserve the size augmentation invariant.

  • Theorems storedSize_rotateLeft_of_wellSized and storedSize_rotateRight_of_wellSized: rotations preserve the cached root size of a well-sized tree.

  • Theorems rankSelect?_rotateLeft and rankSelect?_rotateRight: rotations preserve the ideal rank-selection result.

  • Theorem osSelect?_eq_rankSelect?_of_wellSized: on a well-sized tree, the augmented selector agrees with the ideal rank selector.

  • Theorems osSelect?_rotateLeft_eq_rankSelect?_of_wellSized and osSelect?_rotateRight_eq_rankSelect?_of_wellSized: after a size-preserving rotation, the augmented selector still implements the original ideal rank selector.

  • Theorems rotateLeft_recomputeSizes_wellSized and rotateRight_recomputeSizes_wellSized: recompute-then-rotate produces a well-sized tree from any input tree.

  • Theorems osSelect?_rotateLeft_recomputeSizes_eq_rankSelect? and osSelect?_rotateRight_recomputeSizes_eq_rankSelect?: recompute-then- rotate preserves the augmented selector's agreement with the original ideal rank selector.

Size augmentation through executable red-black insertion

The second half of this file closes the "stored-field refinement" gap by threading the size augmentation through an executable red-black insertion. The augmented tree OSRBTree caches both a node colour and a subtree size. Its Okasaki-style balanceLeft/balanceRight/insertFixup/ insert mirror the Chapter 13 red-black operations node-for-node, but every reconstructed node is built by the smart constructor mk, which recomputes the cached size from its children. Two bridges connect this to the existing Chapter 13 development and to the ideal order-statistic semantics:

  • Theorem OSRBTree.wellSized_insert: the augmentation invariant survives through balancing — inserting into a well-sized tree yields a well-sized tree, so every cached size field is correct after the red-black rebalancing (CLRS 14.1 maintained through RB-INSERT).

  • Theorem OSRBTree.osSelect?_insert_eq_rankSelect?: after insertion the augmented (cached-size) selector still agrees with the ideal recomputed-size rank selector.

  • Theorem OSRBTree.storedSize_insert: the cached root size is correct after insertion.

  • Theorem OSRBTree.toRB_insert: erasing the size field commutes with insertion — the augmented insert refines the executable Chapter 13 RBTree.insert exactly. Through this refinement, Chapter 13's shape, membership, and height theorems transfer to the augmented tree (Theorems OSRBTree.redBlackShape_toRB_insert and OSRBTree.mem_keys_insert).

Deletion is threaded through the augmentation as well: the smart-constructor mirrors OSRBTree.baldL/baldR/splitMin/join/ del/delete refine Chapter 13's executable deletion (OSRBTree.toRB_delete), so Chapter 13's shape and membership theorems transfer (OSRBTree.redBlackShape_toRB_delete, OSRBTree.mem_keys_delete), and OSRBTree.wellSized_delete shows the size invariant survives deletion (CLRS 14.1 maintained through RB-DELETE).

Legacy-layer boundary:

  • This file supplies OS-SELECT and the augmentation-preserving red-black update spine. The canonical fourth-edition §17.1 module adds OS-RANK, proves cached and ideal ranks agree, and connects its cost to the logarithmic red-black height bound.

  • A single combined BST/red-black/size predicate and RAM-level update costs are optional refinements; the current transfer theorems expose those properties separately.

namespace CLRSnamespace Chapter14
Augmented tree model

A binary tree whose internal nodes cache their subtree size.

inductive OSTree where | empty : OSTree | node : OSTree → Nat → Nat → OSTree → OSTree deriving Repr, DecidableEq
namespace OSTree

Mathematical inorder traversal of the keys, ignoring cached sizes.

def keys : OSTree → List Nat | empty => [] | node left key _size right => keys left ++ [key] ++ keys right

The cached size stored at the root. Empty trees have cached size zero.

def storedSize : OSTree → Nat | empty => 0 | node _left _key size _right => size

The mathematical size obtained by recursively counting nodes.

def realSize : OSTree → Nat | empty => 0 | node left _key _size right => realSize left + realSize right + 1

Every cached size field agrees with the mathematical subtree size.

def WellSized : OSTree → Prop | empty => True | node left _key size right => WellSized left ∧ WellSized right ∧ size = realSize left + realSize right + 1

Recompute every cached size field from the children upward.

def recomputeSizes : OSTree → OSTree | empty => empty | node left key _size right => let left' := recomputeSizes left let right' := recomputeSizes right node left' key (realSize left' + realSize right' + 1) right'
Local rotations

Left rotation with local size recomputation. If the right child is empty, the tree is left unchanged.

def rotateLeft : OSTree → OSTree | node a x _ (node b y _ c) => let left' := node a x (realSize a + realSize b + 1) b node left' y (realSize left' + realSize c + 1) c | t => t

Right rotation with local size recomputation. If the left child is empty, the tree is left unchanged.

def rotateRight : OSTree → OSTree | node (node a x _ b) y _ c => let right' := node b y (realSize b + realSize c + 1) c node a x (realSize a + realSize right' + 1) right' | t => t
Selectors

The ideal rank selector, using mathematically recomputed subtree sizes. Ranks are zero-based: rank zero returns the first inorder key.

def rankSelect? : OSTree → Nat → Option Nat | empty, _ => none | node left key _size right, i => if i < realSize left then rankSelect? left i else if i = realSize left then some key else rankSelect? right (i - realSize left - 1)

The augmented order-statistic selector, using cached subtree sizes rather than recomputing them. The main theorem states that this agrees with rankSelect? whenever the cached fields are well-sized.

def osSelect? : OSTree → Nat → Option Nat | empty, _ => none | node left key _size right, i => if i < storedSize left then osSelect? left i else if i = storedSize left then some key else osSelect? right (i - storedSize left - 1)
Augmentation correctness

A well-sized tree has a correct root size field.

theorem storedSize_eq_realSize_of_wellSized {t : OSTree} (h : WellSized t) : storedSize t = realSize t := by cases t with | empty => rfl | node left key size right => exact h.2.2

Recomputing cached size fields preserves the inorder key sequence.

theorem keys_recomputeSizes (t : OSTree) : keys (recomputeSizes t) = keys t := by induction t with | empty => rfl | node left key size right ihLeft ihRight => simp [recomputeSizes, keys, ihLeft, ihRight]

Recomputing cached size fields preserves the mathematical subtree size.

theorem realSize_recomputeSizes (t : OSTree) : realSize (recomputeSizes t) = realSize t := by induction t with | empty => rfl | node left key size right ihLeft ihRight => simp [recomputeSizes, realSize, ihLeft, ihRight]

Recomputing cached size fields establishes the size augmentation invariant.

theorem recomputeSizes_wellSized (t : OSTree) : WellSized (recomputeSizes t) := by induction t with | empty => trivial | node left key size right ihLeft ihRight => simp [recomputeSizes, WellSized, ihLeft, ihRight]
Rotation correctness for the size augmentation

Left rotation preserves the inorder key sequence.

theorem keys_rotateLeft (t : OSTree) : keys (rotateLeft t) = keys t := by cases t with | empty => rfl | node a x sx right => cases right with | empty => rfl | node b y sy c => simp [rotateLeft, keys, List.append_assoc]

Right rotation preserves the inorder key sequence.

theorem keys_rotateRight (t : OSTree) : keys (rotateRight t) = keys t := by cases t with | empty => rfl | node left y sy c => cases left with | empty => rfl | node a x sx b => simp [rotateRight, keys, List.append_assoc]

Left rotation preserves the mathematical subtree size.

theorem realSize_rotateLeft (t : OSTree) : realSize (rotateLeft t) = realSize t := by cases t with | empty => rfl | node a x sx right => cases right with | empty => rfl | node b y sy c => simp [rotateLeft, realSize] omega

Right rotation preserves the mathematical subtree size.

theorem realSize_rotateRight (t : OSTree) : realSize (rotateRight t) = realSize t := by cases t with | empty => rfl | node left y sy c => cases left with | empty => rfl | node a x sx b => simp [rotateRight, realSize] omega

Left rotation preserves the cached root size of a well-sized tree.

theorem storedSize_rotateLeft_of_wellSized {t : OSTree} (h : WellSized t) : storedSize (rotateLeft t) = storedSize t := by cases t with | empty => rfl | node a x sx right => cases right with | empty => rfl | node b y sy c => rcases h with ⟨_ha, _hRight, hSize⟩ simp [rotateLeft, storedSize, realSize] at hSize ⊢ omega

Right rotation preserves the cached root size of a well-sized tree.

theorem storedSize_rotateRight_of_wellSized {t : OSTree} (h : WellSized t) : storedSize (rotateRight t) = storedSize t := by cases t with | empty => rfl | node left y sy c => cases left with | empty => rfl | node a x sx b => rcases h with ⟨_hLeft, _hc, hSize⟩ simp [rotateRight, storedSize, realSize] at hSize ⊢ omega

Left rotation preserves ideal rank selection.

theorem rankSelect?_rotateLeft (t : OSTree) (i : Nat) : rankSelect? (rotateLeft t) i = rankSelect? t i := by cases t with | empty => rfl | node a x sx right => cases right with | empty => rfl | node b y sy c => by_cases hiA : i < realSize a · have hiLeft : i < realSize a + realSize b + 1 := by omega simp [rotateLeft, rankSelect?, realSize, hiA, hiLeft] · by_cases hiEqA : i = realSize a · have hiLeft : i < realSize a + realSize b + 1 := by omega simp [rotateLeft, rankSelect?, realSize, hiEqA] · by_cases hiLeft : i < realSize a + realSize b + 1 · have hjLt : i - realSize a - 1 < realSize b := by omega simp [rotateLeft, rankSelect?, realSize, hiA, hiEqA, hiLeft, hjLt] · have hjNotLt : ¬ i - realSize a - 1 < realSize b := by omega by_cases hjEq : i - realSize a - 1 = realSize b · have hiEqLeft : i = realSize a + realSize b + 1 := by omega subst i have hNotLtA : ¬ realSize a + realSize b + 1 < realSize a := by omega have hNotEqA : ¬ realSize a + realSize b + 1 = realSize a := by omega have hEqB : realSize a + realSize b + 1 - realSize a - 1 = realSize b := by omega simp [rotateLeft, rankSelect?, realSize, hNotLtA, hNotEqA, hEqB] · have hiNeLeft : i ≠ realSize a + realSize b + 1 := by omega have hIndex : i - (realSize a + realSize b + 1) - 1 = i - realSize a - 1 - realSize b - 1 := by omega simp [rotateLeft, rankSelect?, realSize, hiA, hiEqA, hiLeft, hjNotLt, hjEq, hiNeLeft, hIndex]

Right rotation preserves ideal rank selection.

theorem rankSelect?_rotateRight (t : OSTree) (i : Nat) : rankSelect? (rotateRight t) i = rankSelect? t i := by cases t with | empty => rfl | node left y sy c => cases left with | empty => rfl | node a x sx b => by_cases hiA : i < realSize a · have hiLeft : i < realSize a + realSize b + 1 := by omega simp [rotateRight, rankSelect?, realSize, hiA, hiLeft] · by_cases hiEqA : i = realSize a · have hiLeft : i < realSize a + realSize b + 1 := by omega simp [rotateRight, rankSelect?, realSize, hiEqA] · by_cases hiLeft : i < realSize a + realSize b + 1 · have hjLt : i - realSize a - 1 < realSize b := by omega simp [rotateRight, rankSelect?, realSize, hiA, hiEqA, hiLeft, hjLt] · have hjNotLt : ¬ i - realSize a - 1 < realSize b := by omega by_cases hjEq : i - realSize a - 1 = realSize b · have hiEqLeft : i = realSize a + realSize b + 1 := by omega subst i have hNotLtA : ¬ realSize a + realSize b + 1 < realSize a := by omega have hNotEqA : ¬ realSize a + realSize b + 1 = realSize a := by omega have hEqB : realSize a + realSize b + 1 - realSize a - 1 = realSize b := by omega simp [rotateRight, rankSelect?, realSize, hNotLtA, hNotEqA, hEqB] · have hiNeLeft : i ≠ realSize a + realSize b + 1 := by omega have hIndex : i - (realSize a + realSize b + 1) - 1 = i - realSize a - 1 - realSize b - 1 := by omega simp [rotateRight, rankSelect?, realSize, hiA, hiEqA, hiLeft, hjNotLt, hjEq, hiNeLeft, hIndex]

Left rotation with local size recomputation preserves WellSized.

theorem rotateLeft_wellSized {t : OSTree} (h : WellSized t) : WellSized (rotateLeft t) := by cases t with | empty => trivial | node a x sx right => cases right with | empty => simpa [rotateLeft] using h | node b y sy c => rcases h with ⟨ha, hRight, _hSize⟩ rcases hRight with ⟨hb, hc, _hRightSize⟩ simp [rotateLeft, WellSized, realSize, ha, hb, hc]

Right rotation with local size recomputation preserves WellSized.

theorem rotateRight_wellSized {t : OSTree} (h : WellSized t) : WellSized (rotateRight t) := by cases t with | empty => trivial | node left y sy c => cases left with | empty => simpa [rotateRight] using h | node a x sx b => rcases h with ⟨hLeft, hc, _hSize⟩ rcases hLeft with ⟨ha, hb, _hLeftSize⟩ simp [rotateRight, WellSized, realSize, ha, hb, hc]

The augmented selector agrees with the ideal selector on well-sized trees.

theorem osSelect?_eq_rankSelect?_of_wellSized {t : OSTree} {i : Nat} (h : WellSized t) : osSelect? t i = rankSelect? t i := by induction t generalizing i with | empty => rfl | node left key size right ihLeft ihRight => rcases h with ⟨hLeft, hRight, hSize⟩ have hLeftSize : storedSize left = realSize left := storedSize_eq_realSize_of_wellSized hLeft by_cases hlt : i < realSize left · simp [osSelect?, rankSelect?, hLeftSize, hlt, ihLeft hLeft] · by_cases heq : i = realSize left · simp [osSelect?, rankSelect?, hLeftSize, heq] · simp [osSelect?, rankSelect?, hLeftSize, hlt, heq, ihRight hRight]

Recomputing size fields preserves the ideal rank selector.

theorem rankSelect?_recomputeSizes (t : OSTree) (i : Nat) : rankSelect? (recomputeSizes t) i = rankSelect? t i := by induction t generalizing i with | empty => rfl | node left key size right ihLeft ihRight => have hLeftSize : realSize (recomputeSizes left) = realSize left := realSize_recomputeSizes left by_cases hlt : i < realSize left · simp [recomputeSizes, rankSelect?, hLeftSize, hlt, ihLeft] · by_cases heq : i = realSize left · simp [recomputeSizes, rankSelect?, hLeftSize, heq] · simp [recomputeSizes, rankSelect?, hLeftSize, hlt, heq, ihRight]

After a size-preserving left rotation, the augmented selector still implements the original ideal rank selector.

theorem osSelect?_rotateLeft_eq_rankSelect?_of_wellSized {t : OSTree} {i : Nat} (h : WellSized t) : osSelect? (rotateLeft t) i = rankSelect? t i := by calc osSelect? (rotateLeft t) i = rankSelect? (rotateLeft t) i := osSelect?_eq_rankSelect?_of_wellSized (rotateLeft_wellSized h) _ = rankSelect? t i := rankSelect?_rotateLeft t i

After a size-preserving right rotation, the augmented selector still implements the original ideal rank selector.

theorem osSelect?_rotateRight_eq_rankSelect?_of_wellSized {t : OSTree} {i : Nat} (h : WellSized t) : osSelect? (rotateRight t) i = rankSelect? t i := by calc osSelect? (rotateRight t) i = rankSelect? (rotateRight t) i := osSelect?_eq_rankSelect?_of_wellSized (rotateRight_wellSized h) _ = rankSelect? t i := rankSelect?_rotateRight t i

Recomputing size fields makes the augmented selector agree with the ideal selector without requiring an external invariant proof.

theorem osSelect?_recomputeSizes_eq_rankSelect? (t : OSTree) (i : Nat) : osSelect? (recomputeSizes t) i = rankSelect? (recomputeSizes t) i := by exact osSelect?_eq_rankSelect?_of_wellSized (recomputeSizes_wellSized t)

Recomputing size fields and then rotating left produces a well-sized tree.

theorem rotateLeft_recomputeSizes_wellSized (t : OSTree) : WellSized (rotateLeft (recomputeSizes t)) := by exact rotateLeft_wellSized (recomputeSizes_wellSized t)

Recomputing size fields and then rotating right produces a well-sized tree.

theorem rotateRight_recomputeSizes_wellSized (t : OSTree) : WellSized (rotateRight (recomputeSizes t)) := by exact rotateRight_wellSized (recomputeSizes_wellSized t)

After recomputing size fields and rotating left, the augmented selector still implements the original ideal rank selector.

theorem osSelect?_rotateLeft_recomputeSizes_eq_rankSelect? (t : OSTree) (i : Nat) : osSelect? (rotateLeft (recomputeSizes t)) i = rankSelect? t i := by calc osSelect? (rotateLeft (recomputeSizes t)) i = rankSelect? (recomputeSizes t) i := osSelect?_rotateLeft_eq_rankSelect?_of_wellSized (recomputeSizes_wellSized t) _ = rankSelect? t i := rankSelect?_recomputeSizes t i

After recomputing size fields and rotating right, the augmented selector still implements the original ideal rank selector.

theorem osSelect?_rotateRight_recomputeSizes_eq_rankSelect? (t : OSTree) (i : Nat) : osSelect? (rotateRight (recomputeSizes t)) i = rankSelect? t i := by calc osSelect? (rotateRight (recomputeSizes t)) i = rankSelect? (recomputeSizes t) i := osSelect?_rotateRight_eq_rankSelect?_of_wellSized (recomputeSizes_wellSized t) _ = rankSelect? t i := rankSelect?_recomputeSizes t i
end OSTree
Size augmentation through executable red-black insertion

This section threads the size augmentation through an executable red-black insertion. The augmented tree OSRBTree caches both a node colour (reusing CLRS.Chapter13.Color) and a subtree size. The red-black operations are the Okasaki-style ones from Chapter 13, but every reconstructed node is built by the smart constructor OSRBTree.mk, which recomputes the cached size from its children. The headline result OSRBTree.wellSized_insert shows that the size augmentation invariant survives through balancing.

open Chapter13 (Color RBTree)

An augmented red-black tree: every internal node caches a colour and a subtree size. This is the red-black refinement of OSTree, adding the colour field needed to run the Chapter 13 insertion balancer.

inductive OSRBTree where | empty : OSRBTree | node : Color → OSRBTree → Nat → Nat → OSRBTree → OSRBTree deriving Repr, DecidableEq
namespace OSRBTree

Inorder traversal of the keys, ignoring colours and cached sizes.

def keys : OSRBTree → List Nat | empty => [] | node _ left key _ right => keys left ++ [key] ++ keys right

The cached size stored at the root. Empty trees have cached size zero.

def storedSize : OSRBTree → Nat | empty => 0 | node _ _ _ size _ => size

The mathematical size obtained by recursively counting nodes.

def realSize : OSRBTree → Nat | empty => 0 | node _ left _ _ right => realSize left + realSize right + 1

Every cached size field agrees with the mathematical subtree size.

def WellSized : OSRBTree → Prop | empty => True | node _ left _ size right => WellSized left ∧ WellSized right ∧ size = realSize left + realSize right + 1

Smart constructor that recomputes the cached size from the children. Every node produced by the red-black operations below is built with mk, which is why the size augmentation invariant is preserved automatically.

def mk (c : Color) (l : OSRBTree) (k : Nat) (r : OSRBTree) : OSRBTree := node c l k (realSize l + realSize r + 1) r

Erase the cached size field, projecting onto the Chapter 13 red-black tree.

def toRB : OSRBTree → RBTree | empty => RBTree.empty | node c l k _ r => RBTree.node c (toRB l) k (toRB r)
Basic facts about the smart constructor and the erasure

mk recomputes the mathematical subtree size correctly.

theorem realSize_mk (c : Color) (l : OSRBTree) (k : Nat) (r : OSRBTree) : realSize (mk c l k r) = realSize l + realSize r + 1 := rfl

mk preserves the inorder key sequence.

theorem keys_mk (c : Color) (l : OSRBTree) (k : Nat) (r : OSRBTree) : keys (mk c l k r) = keys l ++ [k] ++ keys r := rfl

The cached root size of a mk node is the recomputed size.

theorem storedSize_mk (c : Color) (l : OSRBTree) (k : Nat) (r : OSRBTree) : storedSize (mk c l k r) = realSize l + realSize r + 1 := rfl

A mk node is well-sized whenever both children are.

theorem wellSized_mk {c : Color} {l : OSRBTree} {k : Nat} {r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (mk c l k r) := ⟨hl, hr, rfl⟩

Erasing the cached size of a mk node forgets the size only.

theorem toRB_mk (c : Color) (l : OSRBTree) (k : Nat) (r : OSRBTree) : toRB (mk c l k r) = RBTree.node c (toRB l) k (toRB r) := rfl

A well-sized tree has a correct root size field.

theorem storedSize_eq_realSize_of_wellSized {t : OSRBTree} (h : WellSized t) : storedSize t = realSize t := by cases t with | empty => rfl | node c l k s r => exact h.2.2

Erasure relates keys membership to Chapter 13 tree membership.

theorem inTree_toRB (y : Nat) (t : OSRBTree) : RBTree.InTree y (toRB t) ↔ y ∈ keys t := by induction t with | empty => simp [toRB, RBTree.InTree, keys] | node c l k s r ihl ihr => simp only [toRB, RBTree.InTree, keys, List.append_assoc, List.singleton_append, List.mem_append, List.mem_cons, ihl, ihr] tauto
Selectors

The ideal rank selector using mathematically recomputed subtree sizes. Ranks are zero-based.

def rankSelect? : OSRBTree → Nat → Option Nat | empty, _ => none | node _ left key _ right, i => if i < realSize left then rankSelect? left i else if i = realSize left then some key else rankSelect? right (i - realSize left - 1)

The augmented order-statistic selector using cached subtree sizes.

def osSelect? : OSRBTree → Nat → Option Nat | empty, _ => none | node _ left key _ right, i => if i < storedSize left then osSelect? left i else if i = storedSize left then some key else osSelect? right (i - storedSize left - 1)

The augmented selector agrees with the ideal selector on well-sized trees.

theorem osSelect?_eq_rankSelect?_of_wellSized {t : OSRBTree} {i : Nat} (h : WellSized t) : osSelect? t i = rankSelect? t i := by induction t generalizing i with | empty => rfl | node c l k s r ihl ihr => obtain ⟨hLeft, hRight, hSize⟩ := h have hLeftSize : storedSize l = realSize l := storedSize_eq_realSize_of_wellSized hLeft by_cases hlt : i < realSize l · simp [osSelect?, rankSelect?, hLeftSize, hlt, ihl hLeft] · by_cases heq : i = realSize l · simp [osSelect?, rankSelect?, hLeftSize, heq] · simp [osSelect?, rankSelect?, hLeftSize, hlt, heq, ihr hRight]
Executable red-black operations with size recomputation

Repaint the root black, keeping the cached size fields.

def repaintBlack : OSRBTree → OSRBTree | empty => empty | node _ l k s r => node Color.black l k s r

Okasaki-style rebalance after insertion on the left child, recomputing sizes. Mirrors CLRS.Chapter13.RBTree.balanceLeft.

def balanceLeft (l : OSRBTree) (y : Nat) (r : OSRBTree) : OSRBTree := match l with | node Color.red (node Color.red a w _ b) x _ c => mk Color.red (mk Color.black a w b) x (mk Color.black c y r) | node Color.red a w _ (node Color.red b x _ c) => mk Color.red (mk Color.black a w b) x (mk Color.black c y r) | _ => mk Color.black l y r

Okasaki-style rebalance after insertion on the right child, recomputing sizes. Mirrors CLRS.Chapter13.RBTree.balanceRight.

def balanceRight (l : OSRBTree) (y : Nat) (r : OSRBTree) : OSRBTree := match r with | node Color.red (node Color.red b x _ c) y' _ d => mk Color.red (mk Color.black l y b) x (mk Color.black c y' d) | node Color.red b x _ (node Color.red c y' _ d) => mk Color.red (mk Color.black l y b) x (mk Color.black c y' d) | _ => mk Color.black l y r

Insertion fixup: recurse down the tree and rebalance on the way up.

def insertFixup (x : Nat) : OSRBTree → OSRBTree | empty => mk Color.red empty x empty | node c l y s r => if x < y then if c = Color.black then balanceLeft (insertFixup x l) y r else mk Color.red (insertFixup x l) y r else if x > y then if c = Color.black then balanceRight l y (insertFixup x r) else mk Color.red l y (insertFixup x r) else node c l y s r

Insert a key into an augmented red-black tree and repaint the root black.

def insert (x : Nat) (t : OSRBTree) : OSRBTree := repaintBlack (insertFixup x t)
The augmentation invariant survives balancing

Repainting the root black preserves the size augmentation invariant.

theorem wellSized_repaintBlack {t : OSRBTree} (h : WellSized t) : WellSized (repaintBlack t) := by cases t with | empty => trivial | node c l k s r => exact ⟨h.1, h.2.1, h.2.2⟩

balanceLeft preserves the size augmentation invariant.

theorem wellSized_balanceLeft {l : OSRBTree} {y : Nat} {r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (balanceLeft l y r) := by unfold balanceLeft split · obtain ⟨⟨ha, hb, _⟩, hc, _⟩ := hl exact wellSized_mk (wellSized_mk ha hb) (wellSized_mk hc hr) · obtain ⟨ha, ⟨hb, hc, _⟩, _⟩ := hl exact wellSized_mk (wellSized_mk ha hb) (wellSized_mk hc hr) · exact wellSized_mk hl hr

balanceRight preserves the size augmentation invariant.

theorem wellSized_balanceRight {l : OSRBTree} {y : Nat} {r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (balanceRight l y r) := by unfold balanceRight split · obtain ⟨⟨hb, hc, _⟩, hd, _⟩ := hr exact wellSized_mk (wellSized_mk hl hb) (wellSized_mk hc hd) · obtain ⟨hb, ⟨hc, hd, _⟩, _⟩ := hr exact wellSized_mk (wellSized_mk hl hb) (wellSized_mk hc hd) · exact wellSized_mk hl hr

insertFixup preserves the size augmentation invariant.

theorem wellSized_insertFixup (x : Nat) {t : OSRBTree} (h : WellSized t) : WellSized (insertFixup x t) := by induction t with | empty => simp only [insertFixup] exact wellSized_mk (by trivial) (by trivial) | node c l y s r ihl ihr => have hl : WellSized l := h.1 have hr : WellSized r := h.2.1 simp only [insertFixup] by_cases h1 : x < y · simp only [h1, if_true] by_cases hc : c = Color.black · simp only [hc, if_true] exact wellSized_balanceLeft (ihl hl) hr · simp only [hc, if_false] exact wellSized_mk (ihl hl) hr · simp only [h1, if_false] by_cases h2 : x > y · simp only [h2, if_true] by_cases hc : c = Color.black · simp only [hc, if_true] exact wellSized_balanceRight hl (ihr hr) · simp only [hc, if_false] exact wellSized_mk hl (ihr hr) · simp only [h2, if_false] exact h

Augmentation invariant through executable insertion (CLRS 14.1 through RB-INSERT). Inserting a key into a well-sized augmented red-black tree produces a well-sized tree: every cached subtree-size field remains correct after the red-black rebalancing (CLRS 14.1 maintained through insert).

theorem wellSized_insert (x : Nat) {t : OSRBTree} (h : WellSized t) : WellSized (insert x t) := by unfold insert exact wellSized_repaintBlack (wellSized_insertFixup x h)

After insertion the cached root size equals the mathematical subtree size.

theorem storedSize_insert (x : Nat) {t : OSRBTree} (h : WellSized t) : storedSize (insert x t) = realSize (insert x t) := storedSize_eq_realSize_of_wellSized (wellSized_insert x h)

After insertion the augmented (cached-size) selector still implements the ideal recomputed-size rank selector.

theorem osSelect?_insert_eq_rankSelect? (x : Nat) {t : OSRBTree} (h : WellSized t) (i : Nat) : osSelect? (insert x t) i = rankSelect? (insert x t) i := osSelect?_eq_rankSelect?_of_wellSized (wellSized_insert x h)
Refinement onto the executable Chapter 13 red-black insertion

Erasing the size field commutes with repainting the root black.

theorem toRB_repaintBlack (t : OSRBTree) : toRB (repaintBlack t) = RBTree.repaintRoot Color.black (toRB t) := by cases t with | empty => rfl | node c l k s r => rfl

Erasing the size field commutes with balanceLeft.

theorem toRB_balanceLeft (l : OSRBTree) (y : Nat) (r : OSRBTree) : toRB (balanceLeft l y r) = RBTree.balanceLeft (toRB l) y (toRB r) := by cases l with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the size field commutes with balanceRight.

theorem toRB_balanceRight (l : OSRBTree) (y : Nat) (r : OSRBTree) : toRB (balanceRight l y r) = RBTree.balanceRight (toRB l) y (toRB r) := by cases r with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the size field commutes with insertFixup.

theorem toRB_insertFixup (x : Nat) (t : OSRBTree) : toRB (insertFixup x t) = RBTree.insertFixup x (toRB t) := by induction t with | empty => rfl | node c l y s r ihl ihr => simp only [insertFixup, RBTree.insertFixup, toRB, apply_ite toRB, toRB_mk, toRB_balanceLeft, toRB_balanceRight, ihl, ihr]

Refinement. The augmented insertion refines the executable Chapter 13 red-black insertion: erasing the cached size fields turns OSRBTree.insert into CLRS.Chapter13.RBTree.insert.

theorem toRB_insert (x : Nat) (t : OSRBTree) : toRB (insert x t) = RBTree.insert x (toRB t) := by unfold insert RBTree.insert rw [toRB_repaintBlack, toRB_insertFixup]

Through the refinement, the Chapter 13 red-black shape invariant is maintained by the augmented insertion.

theorem redBlackShape_toRB_insert (x : Nat) {t : OSRBTree} (h : RBTree.RedBlackShape (toRB t)) : RBTree.RedBlackShape (toRB (insert x t)) := by rw [toRB_insert] exact RBTree.redBlackShape_insert h

Through the refinement, insertion preserves membership (as an inorder key).

theorem mem_keys_insert (x y : Nat) (t : OSRBTree) : y ∈ keys (insert x t) ↔ y = x ∨ y ∈ keys t := by simp only [← inTree_toRB, toRB_insert, RBTree.inTree_insert_iff]
Executable red-black deletion with size recomputation

The deletion machinery of Chapter 13, mirrored node-for-node with the smart constructor mk so that the size invariant is preserved automatically. Every definition below refines its Chapter 13 counterpart exactly once the cached size field is erased (see the refinement subsection further down).

Repaint the root with an arbitrary colour, keeping the cached size fields.

def repaintRoot (c : Color) : OSRBTree → OSRBTree | empty => empty | node _ l k s r => node c l k s r

Boolean black-root test. Mirrors CLRS.Chapter13.RBTree.rootBlack.

def rootBlack : OSRBTree → Bool | empty => true | node c _ _ _ _ => c = Color.black

Deletion re-balancer for a black-deficient left child, recomputing sizes. Mirrors CLRS.Chapter13.RBTree.baldL.

def baldL : OSRBTree → Nat → OSRBTree → OSRBTree | node Color.red a x _ b, k, r => mk Color.red (mk Color.black a x b) k r | l, k, node Color.black c y _ d => balanceRight l k (mk Color.red c y d) | l, k, node Color.red (node Color.black c y _ d) z _ e => mk Color.red (mk Color.black l k c) y (balanceRight d z (repaintRoot Color.red e)) | l, k, r => mk Color.red l k r

Deletion re-balancer for a black-deficient right child, recomputing sizes. Mirrors CLRS.Chapter13.RBTree.baldR.

def baldR : OSRBTree → Nat → OSRBTree → OSRBTree | l, k, node Color.red c y _ d => mk Color.red l k (mk Color.black c y d) | node Color.black a x _ b, k, r => balanceLeft (mk Color.red a x b) k r | node Color.red a x _ (node Color.black c y _ d), k, r => mk Color.red (balanceLeft (repaintRoot Color.red a) x c) y (mk Color.black d k r) | l, k, r => mk Color.red l k r

Find and remove the minimum key from a non-empty tree, recomputing sizes on the way back up. Mirrors CLRS.Chapter13.RBTree.splitMin.

def splitMin : OSRBTree → Nat × OSRBTree | empty => (0, empty) -- unreachable on valid inputs | node _ empty k _ r => (k, r) | node _ l k _ r => let (m, l') := splitMin l if rootBlack l then (m, baldL l' k r) else (m, mk Color.red l' k r)

Merge two trees into one, used when deleting a node with two children. Mirrors CLRS.Chapter13.RBTree.join.

def join (l r : OSRBTree) : OSRBTree := if r = empty then l else if l = empty then r else let (m, r') := splitMin r if rootBlack r then baldR l m r' else mk Color.red l m r'

Recursive deletion from an augmented red-black tree, recomputing sizes. Mirrors CLRS.Chapter13.RBTree.del.

def del (x : Nat) : OSRBTree → OSRBTree | empty => empty | node _c l y _s r => if x < y then if rootBlack l then baldL (del x l) y r else mk Color.red (del x l) y r else if x > y then if rootBlack r then baldR l y (del x r) else mk Color.red l y (del x r) else join l r

Delete a key from an augmented red-black tree and repaint the root black. Mirrors CLRS.Chapter13.RBTree.delete.

def delete (x : Nat) (t : OSRBTree) : OSRBTree := repaintBlack (del x t)
The augmentation invariant survives deletion

Repainting the root preserves the size augmentation invariant.

theorem wellSized_repaintRoot (c : Color) {t : OSRBTree} (h : WellSized t) : WellSized (repaintRoot c t) := by cases t with | empty => trivial | node c' l k s r => exact ⟨h.1, h.2.1, h.2.2⟩

baldL preserves the size augmentation invariant.

theorem wellSized_baldL {l : OSRBTree} {k : Nat} {r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (baldL l k r) := by cases l with | empty => cases r with | empty => exact wellSized_mk hl hr | node rc rl rk rs rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellSized_balanceRight hl (wellSized_mk hc hd) | red => cases rl with | empty => exact wellSized_mk hl hr | node rlc rll rlk rls rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellSized_mk (wellSized_mk hl hc) (wellSized_balanceRight hd (wellSized_repaintRoot _ he)) | red => exact wellSized_mk hl hr | node lc ll lk ls lr => cases lc with | red => obtain ⟨ha, hb, _⟩ := hl exact wellSized_mk (wellSized_mk ha hb) hr | black => cases r with | empty => exact wellSized_mk hl hr | node rc rl rk rs rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellSized_balanceRight hl (wellSized_mk hc hd) | red => cases rl with | empty => exact wellSized_mk hl hr | node rlc rll rlk rls rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellSized_mk (wellSized_mk hl hc) (wellSized_balanceRight hd (wellSized_repaintRoot _ he)) | red => exact wellSized_mk hl hr

baldR preserves the size augmentation invariant.

theorem wellSized_baldR {l : OSRBTree} {k : Nat} {r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (baldR l k r) := by cases r with | empty => cases l with | empty => exact wellSized_mk hl hr | node lc ll lk ls lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellSized_balanceLeft (wellSized_mk ha hb) hr | red => cases lr with | empty => exact wellSized_mk hl hr | node lrc lrl lrk lrs lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellSized_mk (wellSized_balanceLeft (wellSized_repaintRoot _ ha) hc) (wellSized_mk hd hr) | red => exact wellSized_mk hl hr | node rc rl rk rs rr => cases rc with | red => obtain ⟨hc, hd, _⟩ := hr exact wellSized_mk hl (wellSized_mk hc hd) | black => cases l with | empty => exact wellSized_mk hl hr | node lc ll lk ls lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellSized_balanceLeft (wellSized_mk ha hb) hr | red => cases lr with | empty => exact wellSized_mk hl hr | node lrc lrl lrk lrs lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellSized_mk (wellSized_balanceLeft (wellSized_repaintRoot _ ha) hc) (wellSized_mk hd hr) | red => exact wellSized_mk hl hr

splitMin preserves the size augmentation invariant.

theorem wellSized_splitMin {t : OSRBTree} (h : WellSized t) : WellSized (splitMin t).2 := by induction t with | empty => trivial | node c l k s r ihl => cases l with | empty => exact h.2.1 | node lc ll lk ls lr => have hws : WellSized (splitMin (node lc ll lk ls lr)).2 := ihl h.1 by_cases hrb : rootBlack (node lc ll lk ls lr) = true · have hsp : (splitMin (node c (node lc ll lk ls lr) k s r)).2 = baldL (splitMin (node lc ll lk ls lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellSized_baldL hws h.2.1 · have hsp : (splitMin (node c (node lc ll lk ls lr) k s r)).2 = mk Color.red (splitMin (node lc ll lk ls lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellSized_mk hws h.2.1

join on two non-empty trees with a black-rooted right tree absorbs the deficit with baldR.

theorem join_eq_baldR {l r : OSRBTree} (hre : r ≠ empty) (hle : l ≠ empty) (hrb : rootBlack r = true) : join l r = baldR l (splitMin r).1 (splitMin r).2 := by simp [join, hre, hle, hrb]

join on two non-empty trees with a red-rooted right tree rebuilds a plain red node (no deficit arises).

theorem join_eq_mk_red {l r : OSRBTree} (hre : r ≠ empty) (hle : l ≠ empty) (hrb : rootBlack r ≠ true) : join l r = mk Color.red l (splitMin r).1 (splitMin r).2 := by simp [join, hre, hle, hrb]

join preserves the size augmentation invariant.

theorem wellSized_join {l r : OSRBTree} (hl : WellSized l) (hr : WellSized r) : WellSized (join l r) := by by_cases hre : r = empty · subst hre exact hl · by_cases hle : l = empty · subst hle have he : join empty r = r := by simp [join, hre] rw [he] exact hr · by_cases hrb : rootBlack r = true · rw [join_eq_baldR hre hle hrb] exact wellSized_baldR hl (wellSized_splitMin hr) · rw [join_eq_mk_red hre hle hrb] exact wellSized_mk hl (wellSized_splitMin hr)

del preserves the size augmentation invariant.

theorem wellSized_del (x : Nat) {t : OSRBTree} (h : WellSized t) : WellSized (del x t) := by induction t with | empty => exact h | node c l y s r ihl ihr => have hl : WellSized l := h.1 have hr : WellSized r := h.2.1 simp only [del] split · split · exact wellSized_baldL (ihl hl) hr · exact wellSized_mk (ihl hl) hr · split · split · exact wellSized_baldR hl (ihr hr) · exact wellSized_mk hl (ihr hr) · exact wellSized_join hl hr

Augmentation invariant through executable deletion (CLRS 14.1 through RB-DELETE). Deleting a key from a well-sized augmented red-black tree produces a well-sized tree: every cached subtree-size field remains correct after the red-black rebalancing (CLRS 14.1 maintained through delete).

theorem wellSized_delete (x : Nat) {t : OSRBTree} (h : WellSized t) : WellSized (delete x t) := by unfold delete exact wellSized_repaintBlack (wellSized_del x h)

After deletion the cached root size equals the mathematical subtree size.

theorem storedSize_delete (x : Nat) {t : OSRBTree} (h : WellSized t) : storedSize (delete x t) = realSize (delete x t) := storedSize_eq_realSize_of_wellSized (wellSized_delete x h)

After deletion the augmented (cached-size) selector still implements the ideal recomputed-size rank selector.

theorem osSelect?_delete_eq_rankSelect? (x : Nat) {t : OSRBTree} (h : WellSized t) (i : Nat) : osSelect? (delete x t) i = rankSelect? (delete x t) i := osSelect?_eq_rankSelect?_of_wellSized (wellSized_delete x h)
Refinement onto the executable Chapter 13 red-black deletion

Erasing the size field commutes with repainting the root.

theorem toRB_repaintRoot (c : Color) (t : OSRBTree) : toRB (repaintRoot c t) = RBTree.repaintRoot c (toRB t) := by cases t with | empty => rfl | node c' l k s r => rfl

Erasing the size field commutes with the black-root test.

theorem toRB_rootBlack (t : OSRBTree) : RBTree.rootBlack (toRB t) = rootBlack t := by cases t with | empty => rfl | node c l k s r => rfl

Erasure hits the empty tree only at the empty tree.

theorem toRB_eq_empty_iff (t : OSRBTree) : toRB t = RBTree.empty ↔ t = empty := by cases t with | empty => simp [toRB] | node c l k s r => simp [toRB]

Erasing the size field commutes with baldL.

theorem toRB_baldL (l : OSRBTree) (k : Nat) (r : OSRBTree) : toRB (baldL l k r) = RBTree.baldL (toRB l) k (toRB r) := by cases l with | empty => cases r with | empty => rfl | node rc rl rk rs rr => cases rc with | black => simp only [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight] | red => cases rl with | empty => rfl | node rlc rll rlk rls rlr => cases rlc with | black => simp only [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot] | red => rfl | node lc ll lk ls lr => cases lc with | red => rfl | black => cases r with | empty => rfl | node rc rl rk rs rr => cases rc with | black => simp only [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight] | red => cases rl with | empty => rfl | node rlc rll rlk rls rlr => cases rlc with | black => simp only [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot] | red => rfl

Erasing the size field commutes with baldR.

theorem toRB_baldR (l : OSRBTree) (k : Nat) (r : OSRBTree) : toRB (baldR l k r) = RBTree.baldR (toRB l) k (toRB r) := by cases r with | empty => cases l with | empty => rfl | node lc ll lk ls lr => cases lc with | black => simp only [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft] | red => cases lr with | empty => rfl | node lrc lrl lrk lrs lrr => cases lrc with | black => simp only [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot] | red => rfl | node rc rl rk rs rr => cases rc with | red => rfl | black => cases l with | empty => rfl | node lc ll lk ls lr => cases lc with | black => simp only [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft] | red => cases lr with | empty => rfl | node lrc lrl lrk lrs lrr => cases lrc with | black => simp only [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot] | red => rfl

Erasing the size field commutes with splitMin.

theorem toRB_splitMin (t : OSRBTree) : (splitMin t).1 = (RBTree.splitMin (toRB t)).1 ∧ toRB (splitMin t).2 = (RBTree.splitMin (toRB t)).2 := by induction t with | empty => exact ⟨rfl, rfl⟩ | node c l k s r ihl => cases l with | empty => exact ⟨rfl, rfl⟩ | node lc ll lk ls lr => obtain ⟨ih1, ih2⟩ := ihl by_cases hrb : rootBlack (node lc ll lk ls lr) = true · have hrb' : RBTree.rootBlack (RBTree.node lc (toRB ll) lk (toRB lr)) = true := hrb refine ⟨?_, ?_⟩ · have h1 : (splitMin (node c (node lc ll lk ls lr) k s r)).1 = (splitMin (node lc ll lk ls lr)).1 := by simp [splitMin, hrb] rw [h1, ih1] simp [RBTree.splitMin, toRB, hrb'] · have h1 : (splitMin (node c (node lc ll lk ls lr) k s r)).2 = baldL (splitMin (node lc ll lk ls lr)).2 k r := by simp [splitMin, hrb] rw [h1, toRB_baldL, ih2] simp [RBTree.splitMin, toRB, hrb'] · have hrb' : RBTree.rootBlack (RBTree.node lc (toRB ll) lk (toRB lr)) ≠ true := hrb refine ⟨?_, ?_⟩ · have h1 : (splitMin (node c (node lc ll lk ls lr) k s r)).1 = (splitMin (node lc ll lk ls lr)).1 := by simp [splitMin, hrb] rw [h1, ih1] simp [RBTree.splitMin, toRB, hrb'] · have h1 : (splitMin (node c (node lc ll lk ls lr) k s r)).2 = mk Color.red (splitMin (node lc ll lk ls lr)).2 k r := by simp [splitMin, hrb] rw [h1, toRB_mk, ih2] simp [RBTree.splitMin, toRB, hrb']

Erasing the size field commutes with join.

theorem toRB_join (l r : OSRBTree) : toRB (join l r) = RBTree.join (toRB l) (toRB r) := by by_cases hr : r = empty · subst hr rfl · by_cases hl : l = empty · subst hl have hre : toRB r ≠ RBTree.empty := fun h => hr ((toRB_eq_empty_iff r).mp h) simp [join, RBTree.join, hr, hre, toRB] · have hre : toRB r ≠ RBTree.empty := fun h => hr ((toRB_eq_empty_iff r).mp h) have hle : toRB l ≠ RBTree.empty := fun h => hl ((toRB_eq_empty_iff l).mp h) by_cases hrb : rootBlack r = true · have hrb' : RBTree.rootBlack (toRB r) = true := by rw [toRB_rootBlack]; exact hrb have e2 : RBTree.join (toRB l) (toRB r) = RBTree.baldR (toRB l) (RBTree.splitMin (toRB r)).1 (RBTree.splitMin (toRB r)).2 := by simp [RBTree.join, hre, hle, hrb'] rw [join_eq_baldR hr hl hrb, e2, toRB_baldR, (toRB_splitMin r).1, (toRB_splitMin r).2] · have hrb' : RBTree.rootBlack (toRB r) ≠ true := by rw [toRB_rootBlack]; exact hrb have e2 : RBTree.join (toRB l) (toRB r) = RBTree.node Color.red (toRB l) (RBTree.splitMin (toRB r)).1 (RBTree.splitMin (toRB r)).2 := by simp [RBTree.join, hre, hle, hrb'] rw [join_eq_mk_red hr hl hrb, e2, toRB_mk, (toRB_splitMin r).1, (toRB_splitMin r).2]

Erasing the size field commutes with del.

theorem toRB_del (x : Nat) (t : OSRBTree) : toRB (del x t) = RBTree.del x (toRB t) := by induction t with | empty => rfl | node c l y s r ihl ihr => simp only [del, RBTree.del, toRB, apply_ite toRB, toRB_mk, toRB_baldL, toRB_baldR, toRB_join, ihl, ihr, ← toRB_rootBlack]

Refinement. The augmented deletion refines the executable Chapter 13 red-black deletion: erasing the cached size fields turns OSRBTree.delete into CLRS.Chapter13.RBTree.delete.

theorem toRB_delete (x : Nat) (t : OSRBTree) : toRB (delete x t) = RBTree.delete x (toRB t) := by unfold delete RBTree.delete rw [toRB_repaintBlack, toRB_del]

Through the refinement, the Chapter 13 red-black shape invariant is maintained by the augmented deletion.

theorem redBlackShape_toRB_delete (x : Nat) {t : OSRBTree} (h : RBTree.RedBlackShape (toRB t)) : RBTree.RedBlackShape (toRB (delete x t)) := by rw [toRB_delete] exact RBTree.redBlackShape_delete h

Through the refinement, deletion preserves membership (as an inorder key); requires the tree to be a binary search tree.

theorem mem_keys_delete (x y : Nat) (t : OSRBTree) (hbst : RBTree.BST (toRB t)) : y ∈ keys (delete x t) ↔ y ∈ keys t ∧ y ≠ x := by simp only [← inTree_toRB, toRB_delete, RBTree.inTree_delete_iff x y (toRB t) hbst]
end OSRBTreeend Chapter14end CLRS
Imports

17.2. How to Augment a Data Structure

The legacy augmentation theorems prove field correctness through functional red-black updates. Their smart constructor recursively recomputes mathematical child augmentations, so that implementation does not justify a local-cost claim.

The Execution companion supplies a distinct cached-field insertion and rotation execution. It reads stored child fields, returns the tree plus actual combine and rotation counters, and refines legacy insertion on WellAugmented inputs. Insertion uses at most 5h + 1 combines and 2h rotations; the red-black height theorem yields logarithmic bounds. Each executed double rotation counts both primitives and their local rebuilds.

The historical augmentationUpdateCost below is only an independent height-based budget. Its theorem does not measure any update by itself. Use AugmentationExecution.insert_maintenanceCost_log_bound for the executed cached insertion. Deletion remains the legacy recomputing implementation, with no attached logarithmic maintenance counter here. Allocation, comparison, and bit-arithmetic internals are outside these field-maintenance counts.

namespace CLRSnamespace Chapter14open CLRS.Chapter13 (RBTree)open AugmentedRBTree (toRB)

Historical height-based analysis budget, not a measured update counter. The cached execution companion proves its own explicit combine/rotation bounds.

def augmentationUpdateCost (c : Nat) {β : Type} (t : AugmentedRBTree Nat β) : Nat := c * (RBTree.height (toRB t) + 1)

The independent height budget is logarithmic on red-black-shaped trees. An actual update bound requires the companion's counted-execution refinement.

theorem augmentation_update_bound (c : Nat) {β : Type} (t : AugmentedRBTree Nat β) (hShape : RBTree.RedBlackShape (toRB t)) : augmentationUpdateCost c t ≤ c * (2 * Nat.log 2 (RBTree.size (toRB t) + 1) + 1) := by simp only [augmentationUpdateCost] have hh := RBTree.height_log_bound (toRB t) hShape exact Nat.mul_le_mul_left c (Nat.add_le_add_right hh 1)
end Chapter14end CLRS

Definitions and proofs

CLRSLean.FourthEdition.Chapter_17.Section_17_2_Augmenting_Data_Structures.Execution

Cached augmentation maintenance during insertion

The execution below reads cached child fields; it never calls realAug. Each constructed node performs and counts one combine. Each successful rotation performs and counts two such constructions and one rotation. Insertion follows one search path and uses these rotation primitives in its balancer. The returned counters belong to that same execution. Comparisons, allocation, bit arithmetic, and reconstruction of persistent tree nodes are not separate units in these augmentation-maintenance counters.

Refinement to the legacy recomputing insertion requires WellAugmented. This module measures insertion and rotations, not deletion.

namespace CLRS.Chapter14.AugmentationExecutionopen CLRS.Chapter13 (Color RBTree)open AugmentedRBTreevariable {α β : Type} [Inhabited β]

Result and the primitive calls made while producing it.

structure Run (α β : Type) where tree : AugmentedRBTree α β combineCalls : Nat rotations : Nat deriving Repr
def pure (t : AugmentedRBTree α β) : Run α β := ⟨t, 0, 0⟩def mapTree (f : AugmentedRBTree α β → AugmentedRBTree α β) (r : Run α β) : Run α β := { r with tree := f r.tree }

One local field recomputation, using only cached child fields.

def make (aug : Augmentation α β) (c : Color) (l : Run α β) (k : α) (r : Run α β) : Run α β := ⟨.node c l.tree k (aug.combine k (storedAug aug l.tree) (storedAug aug r.tree)) r.tree, l.combineCalls + r.combineCalls + 1, l.rotations + r.rotations⟩

The two changed nodes are rebuilt once each. A failed rotation is free.

def rotateLeft (aug : Augmentation α β) (t : Run α β) : Run α β := match t.tree with | .node c a x _ (.node d b y _ e) => let rotated := make aug d (make aug c (pure a) x (pure b)) y (pure e) ⟨rotated.tree, t.combineCalls + rotated.combineCalls, t.rotations + 1⟩ | _ => t
def rotateRight (aug : Augmentation α β) (t : Run α β) : Run α β := match t.tree with | .node c (.node d a x _ b) y _ e => let rotated := make aug d (pure a) x (make aug c (pure b) y (pure e)) ⟨rotated.tree, t.combineCalls + rotated.combineCalls, t.rotations + 1⟩ | _ => t

Color changes do not recompute a field.

def blackenLeft : AugmentedRBTree α β → AugmentedRBTree α β | .empty => .empty | .node c l k a r => .node c (repaintRoot .black l) k a r
def blackenRight : AugmentedRBTree α β → AugmentedRBTree α β | .empty => .empty | .node c l k a r => .node c l k a (repaintRoot .black r)

Single and double rotations are actual calls to the primitives above.

def balanceLeft (aug : Augmentation α β) (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : Run α β := match l with | .node .red (.node .red _ _ _ _) _ _ _ => rotateRight aug (make aug .black (pure (blackenLeft l)) y (pure r)) | .node .red _ _ _ (.node .red _ _ _ _) => rotateRight aug (make aug .black (mapTree blackenLeft (rotateLeft aug (pure l))) y (pure r)) | _ => make aug .black (pure l) y (pure r)
def balanceRight (aug : Augmentation α β) (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : Run α β := match r with | .node .red (.node .red _ _ _ _) _ _ _ => rotateLeft aug (make aug .black (pure l) y (mapTree blackenRight (rotateRight aug (pure r)))) | .node .red _ _ _ (.node .red _ _ _ _) => rotateLeft aug (make aug .black (pure l) y (pure (blackenRight r))) | _ => make aug .black (pure l) y (pure r)

Add calls already performed in a recursive child, without rerunning it.

def after (earlier : Run α β) (next : Run α β) : Run α β := ⟨next.tree, earlier.combineCalls + next.combineCalls, earlier.rotations + next.rotations⟩
def insertFixup (aug : Augmentation α β) (lt : α → α → Bool) (x : α) : AugmentedRBTree α β → Run α β | .empty => make aug .red (pure .empty) x (pure .empty) | .node c l y a r => if lt x y then let child := insertFixup aug lt x l if c = .black then after child (balanceLeft aug child.tree y r) else make aug .red child y (pure r) else if lt y x then let child := insertFixup aug lt x r if c = .black then after child (balanceRight aug l y child.tree) else make aug .red (pure l) y child else pure (.node c l y a r)def insert (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) : Run α β := mapTree repaintBlack (insertFixup aug lt x t)theorem make_refines (aug : Augmentation α β) (c : Color) (l r : Run α β) (k : α) (hl : WellAugmented aug l.tree) (hr : WellAugmented aug r.tree) : (make aug c l k r).tree = mk aug c l.tree k r.tree := by simp only [make, mk, storedAug_eq_realAug_of_wellAugmented aug hl, storedAug_eq_realAug_of_wellAugmented aug hr]private theorem stored_node (aug : Augmentation α β) (c : Color) (l : AugmentedRBTree α β) (k : α) (a : β) (r : AugmentedRBTree α β) : storedAug aug (.node c l k a r) = a := rfltheorem balanceLeft_refines (aug : Augmentation α β) (l r : AugmentedRBTree α β) (k : α) (hl : WellAugmented aug l) (hr : WellAugmented aug r) : (balanceLeft aug l k r).tree = AugmentedRBTree.balanceLeft aug l k r := by unfold balanceLeft split <;> (try simp_all only [WellAugmented]) <;> simp_all [AugmentedRBTree.balanceLeft, rotateLeft, rotateRight, make, mapTree, pure, blackenLeft, repaintRoot, mk, stored_node, realAug, storedAug_eq_realAug_of_wellAugmented]theorem balanceRight_refines (aug : Augmentation α β) (l r : AugmentedRBTree α β) (k : α) (hl : WellAugmented aug l) (hr : WellAugmented aug r) : (balanceRight aug l k r).tree = AugmentedRBTree.balanceRight aug l k r := by unfold balanceRight split <;> (try simp_all only [WellAugmented]) <;> simp_all [AugmentedRBTree.balanceRight, rotateLeft, rotateRight, make, mapTree, pure, blackenRight, repaintRoot, mk, stored_node, realAug, storedAug_eq_realAug_of_wellAugmented] theorem insertFixup_refines (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) (h : WellAugmented aug t) : (insertFixup aug lt x t).tree = AugmentedRBTree.insertFixup aug lt x t := by induction t with | empty => rfl | node c l y a r ihl ihr => have hl := h.1 have hr := h.2.1 have il := ihl hl have ir := ihr hr have wl : WellAugmented aug (insertFixup aug lt x l).tree := by rw [il]; exact wellAugmented_insertFixup aug lt x hl have wr : WellAugmented aug (insertFixup aug lt x r).tree := by rw [ir]; exact wellAugmented_insertFixup aug lt x hr simp only [insertFixup, AugmentedRBTree.insertFixup] split · split · simp only [after] rw [balanceLeft_refines aug _ _ _ wl hr, il] · rw [make_refines aug _ _ _ _ wl hr, il]; rfl · split · split · simp only [after] rw [balanceRight_refines aug _ _ _ hl wr, ir] · rw [make_refines aug _ _ _ _ hl wr, ir]; rfl · rfltheorem insert_refines (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) (h : WellAugmented aug t) : (insert aug lt x t).tree = AugmentedRBTree.insert aug lt x t := by simp only [insert, mapTree, AugmentedRBTree.insert, insertFixup_refines aug lt x t h] theorem insert_wellAugmented (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) (h : WellAugmented aug t) : WellAugmented aug (insert aug lt x t).tree := by rw [insert_refines aug lt x t h] exact wellAugmented_insert aug lt x h

The cached implementation erases to the existing functional RB insertion.

theorem insert_toRB (aug : Augmentation Nat β) (x : Nat) (t : AugmentedRBTree Nat β) (h : WellAugmented aug t) : toRB (insert aug natLt x t).tree = RBTree.insert x (toRB t) := by rw [insert_refines aug natLt x t h, AugmentedRBTree.toRB_insert]

A rotation either returns the original run or adds exactly two combines and one rotation.

theorem rotateLeft_counts (aug : Augmentation α β) (t : Run α β) : rotateLeft aug t = t ∨ ((rotateLeft aug t).combineCalls = t.combineCalls + 2 ∧ (rotateLeft aug t).rotations = t.rotations + 1) := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => exact Or.inl rfl | node c l x v r => cases r with | empty => exact Or.inl rfl | node d b y w e => exact Or.inr ⟨rfl, rfl⟩
theorem rotateRight_counts (aug : Augmentation α β) (t : Run α β) : rotateRight aug t = t ∨ ((rotateRight aug t).combineCalls = t.combineCalls + 2 ∧ (rotateRight aug t).rotations = t.rotations + 1) := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => exact Or.inl rfl | node c l x v r => cases l with | empty => exact Or.inl rfl | node d a x w b => exact Or.inr ⟨rfl, rfl⟩theorem rotateLeft_wellAugmented (aug : Augmentation α β) (t : Run α β) (h : WellAugmented aug t.tree) : WellAugmented aug (rotateLeft aug t).tree := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => trivial | node c l x v r => cases r with | empty => exact h | node d b y w e => simp_all [rotateLeft, make, pure, WellAugmented, realAug, storedAug_eq_realAug_of_wellAugmented]theorem rotateRight_wellAugmented (aug : Augmentation α β) (t : Run α β) (h : WellAugmented aug t.tree) : WellAugmented aug (rotateRight aug t).tree := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => trivial | node c l y v r => cases l with | empty => exact h | node d a x w b => simp_all [rotateRight, make, pure, WellAugmented, realAug, storedAug_eq_realAug_of_wellAugmented]

Erasure identifies each counted rotation with the ordinary RB primitive.

theorem rotateLeft_toRB (aug : Augmentation Nat β) (t : Run Nat β) : toRB (rotateLeft aug t).tree = RBTree.rotateLeft (toRB t.tree) := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => rfl | node c l x v r => cases r <;> rfl
theorem rotateRight_toRB (aug : Augmentation Nat β) (t : Run Nat β) : toRB (rotateRight aug t).tree = RBTree.rotateRight (toRB t.tree) := by rcases t with ⟨t, cc, rr⟩ cases t with | empty => rfl | node c l x v r => cases l <;> rfl

Generic-key structural height; cache values do not affect it.

def height : AugmentedRBTree α β → Nat | .empty => 0 | .node _ l _ _ r => max (height l) (height r) + 1
theorem balanceLeft_counts (aug : Augmentation α β) (l r : AugmentedRBTree α β) (k : α) : (balanceLeft aug l k r).combineCalls ≤ 5 ∧ (balanceLeft aug l k r).rotations ≤ 2 := by unfold balanceLeft split <;> simp [rotateLeft, rotateRight, make, pure, mapTree, blackenLeft, repaintRoot]theorem balanceRight_counts (aug : Augmentation α β) (l r : AugmentedRBTree α β) (k : α) : (balanceRight aug l k r).combineCalls ≤ 5 ∧ (balanceRight aug l k r).rotations ≤ 2 := by unfold balanceRight split <;> simp [rotateLeft, rotateRight, make, pure, mapTree, blackenRight, repaintRoot]theorem insertFixup_counts (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) : (insertFixup aug lt x t).combineCalls ≤ 5 * height t + 1 ∧ (insertFixup aug lt x t).rotations ≤ 2 * height t := by induction t with | empty => simp [insertFixup, make, pure, height] | node c l y a r il ir => have hleft := Nat.le_max_left (height l) (height r) have hright := Nat.le_max_right (height l) (height r) simp only [insertFixup, height] split · split · have hb := balanceLeft_counts aug (insertFixup aug lt x l).tree r y simp only [after] omega · simp only [make, pure] omega · split · split · have hb := balanceRight_counts aug l (insertFixup aug lt x r).tree y simp only [after] omega · simp only [make, pure] omega · simp [pure]theorem insert_counts (aug : Augmentation α β) (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) : (insert aug lt x t).combineCalls ≤ 5 * height t + 1 ∧ (insert aug lt x t).rotations ≤ 2 * height t := insertFixup_counts aug lt x tomit [Inhabited β] in theorem height_eq_toRB (t : AugmentedRBTree Nat β) : height t = RBTree.height (toRB t) := by induction t with | empty => rfl | node c l y a r il ir => simp [height, toRB, RBTree.height, il, ir, Nat.add_comm]

The bound applies to the actual cached execution's combine counter.

theorem insert_combineCalls_log_bound (aug : Augmentation Nat β) (lt : Nat → Nat → Bool) (x : Nat) (t : AugmentedRBTree Nat β) (hs : RBTree.RedBlackShape (toRB t)) : (insert aug lt x t).combineCalls ≤ 10 * Nat.log 2 (RBTree.size (toRB t) + 1) + 1 := by have hc := (insert_counts aug lt x t).1 rw [height_eq_toRB] at hc have hh := RBTree.height_log_bound (toRB t) hs omega
theorem insert_rotations_log_bound (aug : Augmentation Nat β) (lt : Nat → Nat → Bool) (x : Nat) (t : AugmentedRBTree Nat β) (hs : RBTree.RedBlackShape (toRB t)) : (insert aug lt x t).rotations ≤ 4 * Nat.log 2 (RBTree.size (toRB t) + 1) := by have hc := (insert_counts aug lt x t).2 rw [height_eq_toRB] at hc have hh := RBTree.height_log_bound (toRB t) hs omega

Constant charges per counted combine and rotation, on the same run.

def maintenanceCost (combineCharge rotationCharge : Nat) (r : Run α β) : Nat := combineCharge * r.combineCalls + rotationCharge * r.rotations
theorem insert_maintenanceCost_log_bound (aug : Augmentation Nat β) (lt : Nat → Nat → Bool) (x : Nat) (t : AugmentedRBTree Nat β) (hs : RBTree.RedBlackShape (toRB t)) (combineCharge rotationCharge : Nat) : maintenanceCost combineCharge rotationCharge (insert aug lt x t) ≤ combineCharge * (10 * Nat.log 2 (RBTree.size (toRB t) + 1) + 1) + rotationCharge * (4 * Nat.log 2 (RBTree.size (toRB t) + 1)) := by exact Nat.add_le_add (Nat.mul_le_mul_left _ (insert_combineCalls_log_bound aug lt x t hs)) (Nat.mul_le_mul_left _ (insert_rotations_log_bound aug lt x t hs))end CLRS.Chapter14.AugmentationExecution

CLRSLean.Chapter_14.Section_14_3_Interval_Trees

Section 14.3 - Interval trees

This file formalizes the second classic augmentation from CLRS Chapter 14: interval trees. Each node stores an interval and a cached maximum high endpoint of its subtree. The search algorithm uses the cached maximum to prune the left subtree when no overlap is possible there.

To make the pattern reusable, we first define a small generic augmentation framework: an AugmentedTree α β carries a value of type α at each node and a cached augmentation of type β. An Augmentation α β provides the empty default and the local recombination function. The framework proves that recomputing the augmentation from the children preserves the WellAugmented invariant, and that rotations preserve the inorder key sequence and the WellAugmented invariant for augmentations whose combine operator behaves like max (which is the case for the interval-tree instantiation). We then instantiate it to interval trees with the max-high augmentation, and show that the executable intervalSearch? is correct on well-augmented BSTs.

Main results:

  • Generic AugmentedTree lemmas: keys_recompute, realAug_recompute, recompute_wellAugmented, rotateLeft_wellAugmented, rotateRight_wellAugmented.

  • Interval-tree correctness: intervalSearch?_some_overlap and intervalSearch?_none_noOverlap (combined as intervalSearch?_spec).

  • General augmentation theorem (CLRS Theorem 14.1): augmentation_theorem packages that rotations, recomputation, and generic BST insert maintain the WellAugmented invariant and the semantic augmentation for any rotation-invariant augmentation.

  • Size augmentation instance: sizeAug with realAug_sizeAug_eq_length, showing order-statistic size caching is an instance of the same framework.

  • Red-black bridge: rb_augmentation_bridge shows Chapter 13's red-black rotations and root recoloring preserve any rotation-invariant augmentation's value (and the inorder key list), so the augmentation is maintainable through the red-black operations.

  • General augmentation interface: AugmentedRBTree threads an arbitrary Augmentation through an executable red-black insertion; its smart constructor AugmentedRBTree.mk recomputes the cached value, so AugmentedRBTree.wellAugmented_insert shows the invariant survives balancing and AugmentedRBTree.toRB_insert shows the augmentation-erasing projection refines Chapter 13's executable RBTree.insert. Both the size and interval instances are recovered from it (AugmentedRBTree.sizeAug_wellAugmented_insert, AugmentedRBTree.maxHighAug_wellAugmented_insert).

Status: the static interval-search specification, generic local augmentation invariant, and arbitrary cached-field insertion/deletion pipeline are proved. The complete fourth-edition interval-tree and augmentation interfaces remain partial. Missing bridges include the constant-time-combine asymptotic theorem, combined BST/red-black/augmentation preservation, and interval-specific update semantics connecting AugmentedRBTree to the separate static IntervalTree search model. The current low-endpoint-only comparator also needs an equal-low policy before arbitrary intervals can be inserted distinctly.

namespace CLRSnamespace Chapter14
Generic augmented trees

An augmentation schema for a binary tree: a default for the empty tree and a local recombination function.

structure Augmentation (α β : Type) [Inhabited β] where base : β combine : α → β → β → β

Typeclass asserting that an augmentation's combine operation satisfies the rotation-invariance law required for BST rotations to preserve the cached augmentation. The max-high augmentation is the motivating example.

class IsRotationInvariant {α β : Type} [Inhabited β] (aug : Augmentation α β) : Prop where combine_rotate : ∀ (x y : α) (a b c : β), aug.combine y (aug.combine x a b) c = aug.combine x a (aug.combine y b c)

A binary tree whose internal nodes cache an augmentation value.

inductive AugmentedTree (α β : Type) where | empty : AugmentedTree α β | node : AugmentedTree α β → α → β → AugmentedTree α β → AugmentedTree α β deriving Repr, DecidableEq
namespace AugmentedTreevariable {α β : Type} [Inhabited β] (aug : Augmentation α β)

Inorder traversal of the stored values.

@[simp] def keys : AugmentedTree α β → List α | empty => [] | node left key _ right => keys left ++ [key] ++ keys right

The cached augmentation at the root.

@[simp] def storedAug : AugmentedTree α β → β | empty => aug.base | node _ _ a _ => a

The mathematically correct augmentation computed from the children.

@[simp] def realAug : AugmentedTree α β → β | empty => aug.base | node left key _ right => aug.combine key (realAug left) (realAug right)

Every cached augmentation agrees with the mathematically correct one.

@[simp] def WellAugmented : AugmentedTree α β → Prop | empty => True | node left key a right => WellAugmented left ∧ WellAugmented right ∧ a = realAug aug (node left key a right)

Recompute every cached augmentation from the children upward.

@[simp] def recompute : AugmentedTree α β → AugmentedTree α β | empty => empty | node left key _ right => let left' := recompute left let right' := recompute right node left' key (aug.combine key (realAug aug left') (realAug aug right')) right'

Left rotation with local augmentation recomputation.

@[simp] def rotateLeft : AugmentedTree α β → AugmentedTree α β | node a x _ (node b y _ c) => let left' := node a x (aug.combine x (realAug aug a) (realAug aug b)) b node left' y (aug.combine y (realAug aug left') (realAug aug c)) c | t => t

Right rotation with local augmentation recomputation.

@[simp] def rotateRight : AugmentedTree α β → AugmentedTree α β | node (node a x _ b) y _ c => let right' := node b y (aug.combine y (realAug aug b) (realAug aug c)) c node a x (aug.combine x (realAug aug a) (realAug aug right')) right' | t => t

Recomputing preserves the inorder key sequence.

theorem keys_recompute (t : AugmentedTree α β) : keys (recompute aug t) = keys t := by induction t with | empty => rfl | node left key _ right ihLeft ihRight => simp [recompute, keys, ihLeft, ihRight]

Recomputing preserves the mathematical augmentation.

theorem realAug_recompute (t : AugmentedTree α β) : realAug aug (recompute aug t) = realAug aug t := by induction t with | empty => rfl | node left key _ right ihLeft ihRight => simp [recompute, realAug, ihLeft, ihRight]

Recomputing establishes the well-augmented invariant.

theorem recompute_wellAugmented (t : AugmentedTree α β) : WellAugmented aug (recompute aug t) := by induction t with | empty => trivial | node left key _ right ihLeft ihRight => simp [recompute, WellAugmented, realAug_recompute, ihLeft, ihRight]

A well-augmented tree has a correct root augmentation.

theorem storedAug_eq_realAug_of_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : storedAug aug t = realAug aug t := by cases t with | empty => rfl | node left key a right => exact h.2.2

Left rotation preserves the inorder key sequence.

theorem keys_rotateLeft (t : AugmentedTree α β) : keys (rotateLeft aug t) = keys t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => simp [rotateLeft, keys, List.append_assoc]

Right rotation preserves the inorder key sequence.

theorem keys_rotateRight (t : AugmentedTree α β) : keys (rotateRight aug t) = keys t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => simp [rotateRight, keys, List.append_assoc]

Left rotation preserves the mathematical augmentation.

theorem realAug_rotateLeft (t : AugmentedTree α β) [IsRotationInvariant aug] : realAug aug (rotateLeft aug t) = realAug aug t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => simp [rotateLeft, realAug] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves the mathematical augmentation.

theorem realAug_rotateRight (t : AugmentedTree α β) [IsRotationInvariant aug] : realAug aug (rotateRight aug t) = realAug aug t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => simp [rotateRight, realAug] rw [IsRotationInvariant.combine_rotate]

Left rotation preserves the cached root augmentation of a well-augmented tree.

theorem storedAug_rotateLeft_of_wellAugmented {t : AugmentedTree α β} [IsRotationInvariant aug] (h : WellAugmented aug t) : storedAug aug (rotateLeft aug t) = storedAug aug t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => rcases h with ⟨_ha, hRight, hSize⟩ simp [rotateLeft, storedAug, realAug, hSize] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves the cached root augmentation of a well-augmented tree.

theorem storedAug_rotateRight_of_wellAugmented {t : AugmentedTree α β} [IsRotationInvariant aug] (h : WellAugmented aug t) : storedAug aug (rotateRight aug t) = storedAug aug t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => rcases h with ⟨hLeft, _hc, hSize⟩ simp [rotateRight, storedAug, realAug, hSize] rw [IsRotationInvariant.combine_rotate]

Left rotation preserves the well-augmented invariant.

theorem rotateLeft_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (rotateLeft aug t) := by cases t with | empty => exact h | node a x _ right => cases right with | empty => simpa [rotateLeft] using h | node b y _ c => rcases h with ⟨ha, hRight, hSize⟩ rcases hRight with ⟨hb, hc, hRightSize⟩ simp [rotateLeft, WellAugmented, realAug, ha, hb, hc]

Right rotation preserves the well-augmented invariant.

theorem rotateRight_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (rotateRight aug t) := by cases t with | empty => exact h | node left y _ c => cases left with | empty => simpa [rotateRight] using h | node a x _ b => rcases h with ⟨hLeft, hc, hSize⟩ rcases hLeft with ⟨ha, hb, hLeftSize⟩ simp [rotateRight, WellAugmented, realAug, ha, hb, hc]
Generic BST insertion and the general augmentation theorem (CLRS 14.1)

BST insertion by a Boolean comparison lt, recomputing the cached augmentation locally at each node. Generic over any augmentation.

def insert (lt : α → α → Bool) (x : α) : AugmentedTree α β → AugmentedTree α β | empty => node empty x (aug.combine x aug.base aug.base) empty | node l k _ r => if lt x k then let l' := insert lt x l node l' k (aug.combine k (realAug aug l') (realAug aug r)) r else if lt k x then let r' := insert lt x r node l k (aug.combine k (realAug aug l) (realAug aug r')) r' else node l k (aug.combine k (realAug aug l) (realAug aug r)) r

Insertion adds exactly the inserted key to the inorder key multiset.

theorem mem_keys_insert (lt : α → α → Bool) (x y : α) (t : AugmentedTree α β) : y ∈ keys (insert aug lt x t) → y = x ∨ y ∈ keys t := by induction t with | empty => simp [insert, keys] | node l k a r ihl ihr => simp only [insert] split · simp only [keys, List.mem_append, List.mem_singleton] rintro ((h | h) | h) · rcases ihl h with h' | h' <;> tauto · tauto · tauto · split · simp only [keys, List.mem_append, List.mem_singleton] rintro ((h | h) | h) · tauto · tauto · rcases ihr h with h' | h' <;> tauto · simp only [keys, List.mem_append, List.mem_singleton] tauto

Generic insertion preserves the well-augmented invariant.

theorem insert_wellAugmented (lt : α → α → Bool) (x : α) {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (insert aug lt x t) := by induction t with | empty => simp [insert, WellAugmented, realAug] | node l k a r ihl ihr => rcases h with ⟨hl, hr, _ha⟩ simp only [insert] split · exact ⟨ihl hl, hr, by simp [realAug]⟩ · split · exact ⟨hl, ihr hr, by simp [realAug]⟩ · exact ⟨hl, hr, by simp [realAug]⟩

CLRS Theorem 14.1 (maintainability of augmentations). For any locally-computable, rotation-invariant augmentation, every structural primitive used by red-black insertion and deletion — left and right rotation, subtree recomputation, and BST insertion — preserves the inorder key sequence and the mathematical augmentation, and preserves (or re-establishes) the WellAugmented invariant. Hence the augmentation can be maintained through the red-black operations.

theorem augmentation_theorem [IsRotationInvariant aug] : (∀ t : AugmentedTree α β, WellAugmented aug t → WellAugmented aug (rotateLeft aug t)) ∧ (∀ t : AugmentedTree α β, WellAugmented aug t → WellAugmented aug (rotateRight aug t)) ∧ (∀ t : AugmentedTree α β, keys (rotateLeft aug t) = keys t) ∧ (∀ t : AugmentedTree α β, keys (rotateRight aug t) = keys t) ∧ (∀ t : AugmentedTree α β, realAug aug (rotateLeft aug t) = realAug aug t) ∧ (∀ t : AugmentedTree α β, realAug aug (rotateRight aug t) = realAug aug t) ∧ (∀ t : AugmentedTree α β, WellAugmented aug (recompute aug t)) ∧ (∀ (lt : α → α → Bool) (x : α) (t : AugmentedTree α β), WellAugmented aug t → WellAugmented aug (insert aug lt x t)) := ⟨fun _ h => rotateLeft_wellAugmented aug h, fun _ h => rotateRight_wellAugmented aug h, fun t => keys_rotateLeft aug t, fun t => keys_rotateRight aug t, fun t => realAug_rotateLeft aug t, fun t => realAug_rotateRight aug t, fun t => recompute_wellAugmented aug t, fun lt x _ h => insert_wellAugmented aug lt x h⟩
end AugmentedTree
Size augmentation (order-statistic trees as an instance)

The order-statistic augmentation of Section 14.1 — caching each subtree's node count — is an instance of the same generic framework, demonstrating CLRS Theorem 14.1 for a second concrete field alongside interval trees' max-high.

The subtree-size augmentation: the cached value is the number of nodes.

def sizeAug (α : Type) : Augmentation α Nat := ⟨0, fun _ l r => 1 + l + r⟩
instance (α : Type) : IsRotationInvariant (sizeAug α) where combine_rotate x y a b c := by simp [sizeAug]; omega

The size augmentation's mathematical value is exactly the node count.

theorem realAug_sizeAug_eq_length {α : Type} (t : AugmentedTree α Nat) : AugmentedTree.realAug (sizeAug α) t = (AugmentedTree.keys t).length := by induction t with | empty => rfl | node l k a r ihl ihr => simp only [AugmentedTree.realAug] rw [ihl, ihr] simp only [sizeAug, AugmentedTree.keys, List.length_append, List.length_cons, List.length_nil] omega

Interval trees

A closed interval of natural numbers.

def Interval := Nat × Nat
def Interval.low (i : Interval) : Nat := i.1def Interval.high (i : Interval) : Nat := i.2def Interval.overlaps (i j : Interval) : Bool := i.low ≤ j.high && j.low ≤ i.high@[simp] theorem Interval.overlaps_iff {i j : Interval} : Interval.overlaps i j = true ↔ i.low ≤ j.high ∧ j.low ≤ i.high := by simp [Interval.overlaps]

Interval trees are augmented trees whose node value is an interval and whose augmentation is the maximum high endpoint in the subtree.

abbrev IntervalTree := AugmentedTree Interval Nat
def IntervalTree.maxHighAug : Augmentation Interval Nat := ⟨0, fun i l r => max i.high (max l r)⟩instance : IsRotationInvariant IntervalTree.maxHighAug where combine_rotate x y a b c := by simp [IntervalTree.maxHighAug] ac_rflnamespace IntervalTree

Inorder list of intervals.

def keys : IntervalTree → List Interval := AugmentedTree.keys

Cached maximum high endpoint at the root.

def storedMaxHigh : IntervalTree → Nat := AugmentedTree.storedAug maxHighAug

Mathematical maximum high endpoint in the subtree.

def realMaxHigh : IntervalTree → Nat := AugmentedTree.realAug maxHighAug

The max-high augmentation invariant.

def WellAugmented : IntervalTree → Prop := AugmentedTree.WellAugmented maxHighAug

Recompute every cached max-high field.

Left rotation with local max-high recomputation.

Right rotation with local max-high recomputation.

Every interval in the tree has low endpoint at most x.

def allLowLE (t : IntervalTree) (x : Nat) : Prop := ∀ i ∈ keys t, i.low ≤ x

Every interval in the tree has low endpoint at least x.

def allLowGE (t : IntervalTree) (x : Nat) : Prop := ∀ i ∈ keys t, x ≤ i.low

Binary-search-tree ordering by interval low endpoint.

def IsBST : IntervalTree → Prop | AugmentedTree.empty => True | AugmentedTree.node left int _ right => IsBST left ∧ IsBST right ∧ allLowLE left int.low ∧ allLowGE right int.low

Boolean emptiness test for interval trees.

def isEmpty : IntervalTree → Bool | AugmentedTree.empty => true | AugmentedTree.node _ _ _ _ => false

Decision to recurse into the left subtree during interval search.

def goLeft (left : IntervalTree) (q : Interval) : Bool := !isEmpty left && decide (storedMaxHigh left ≥ q.low)

The executable interval-search algorithm from CLRS.

def intervalSearch? : IntervalTree → Interval → Option Interval | AugmentedTree.empty, _ => none | AugmentedTree.node left int _ right, q => if Interval.overlaps int q then some int else if goLeft left q then intervalSearch? left q else intervalSearch? right q

Does the tree contain an interval overlapping the query?

def hasOverlap (t : IntervalTree) (q : Interval) : Prop := ∃ i ∈ keys t, Interval.overlaps i q
end IntervalTreenamespace IntervalTree@[simp] theorem keys_empty : keys AugmentedTree.empty = [] := by simp [keys]@[simp] theorem keys_node {left right : IntervalTree} {int : Interval} {mx : Nat} : keys (AugmentedTree.node left int mx right) = keys left ++ [int] ++ keys right := by simp [keys]@[simp] theorem isEmpty_empty : isEmpty AugmentedTree.empty = true := by rfl@[simp] theorem isEmpty_node {left right : IntervalTree} {int : Interval} {mx : Nat} : isEmpty (AugmentedTree.node left int mx right) = false := by rfl

A true goLeft condition means the left subtree is non-empty and its max-high is at least the query low.

theorem goLeft_true {left : IntervalTree} {q : Interval} (h : goLeft left q = true) : left ≠ AugmentedTree.empty ∧ storedMaxHigh left ≥ q.low := by simp [goLeft, Bool.and_eq_true] at h rcases h with ⟨hne, hmax⟩ constructor · cases left with | empty => simp at hne | node => simp · exact hmax

A false goLeft condition means the left subtree is empty or its max-high is below the query low.

theorem goLeft_false {left : IntervalTree} {q : Interval} (h : goLeft left q = false) : left = AugmentedTree.empty ∨ storedMaxHigh left < q.low := by simp [goLeft] at h cases left with | empty => left; rfl | node left int mx right => right have hmax : decide (storedMaxHigh (AugmentedTree.node left int mx right) ≥ q.low) = false := by simpa using h simpa using hmax

Recomputing cached max-high fields preserves the inorder key sequence.

theorem keys_recompute (t : IntervalTree) : keys (recompute t) = keys t := AugmentedTree.keys_recompute maxHighAug t

Recomputing cached max-high fields preserves the mathematical max-high.

theorem realMaxHigh_recompute (t : IntervalTree) : realMaxHigh (recompute t) = realMaxHigh t := AugmentedTree.realAug_recompute maxHighAug t

Recomputing cached max-high fields establishes the augmentation invariant.

theorem recompute_wellAugmented (t : IntervalTree) : WellAugmented (recompute t) := AugmentedTree.recompute_wellAugmented maxHighAug t

A well-augmented interval tree has a correct root max-high field.

theorem storedMaxHigh_eq_realMaxHigh_of_wellAugmented {t : IntervalTree} (h : WellAugmented t) : storedMaxHigh t = realMaxHigh t := AugmentedTree.storedAug_eq_realAug_of_wellAugmented maxHighAug h

Left rotation preserves the well-augmented invariant.

theorem rotateLeft_wellAugmented {t : IntervalTree} (h : WellAugmented t) : WellAugmented (rotateLeft t) := AugmentedTree.rotateLeft_wellAugmented maxHighAug h

Right rotation preserves the well-augmented invariant.

Membership in keys respects the inorder list membership relation.

@[simp] theorem mem_keys {t : IntervalTree} {i : Interval} : i ∈ keys t ↔ i ∈ AugmentedTree.keys t := by rfl

Every stored high endpoint is bounded by the real max-high.

theorem high_le_realMaxHigh {t : IntervalTree} {i : Interval} (hi : i ∈ keys t) : i.high ≤ realMaxHigh t := by induction t with | empty => simp [keys] at hi | node left int _ right ihLeft ihRight => simp [keys] at hi rcases hi with hi | hi | hi · exact le_trans (ihLeft hi) (by simp [realMaxHigh, AugmentedTree.realAug, maxHighAug]) · simp [hi, realMaxHigh, AugmentedTree.realAug, maxHighAug] · exact le_trans (ihRight hi) (by simp [realMaxHigh, AugmentedTree.realAug, maxHighAug])

If the mathematical max-high of a subtree is below the query low, the subtree contains no overlap.

theorem noOverlap_of_realMaxHigh_lt {t : IntervalTree} {q : Interval} (h : realMaxHigh t < q.low) : ¬ ∃ i ∈ keys t, Interval.overlaps i q := by rintro ⟨i, hi, hov⟩ rw [Interval.overlaps_iff] at hov have hiHigh := high_le_realMaxHigh hi have : q.low ≤ i.high := hov.2 linarith

A tree whose realMaxHigh is at least a positive bound contains a member whose high is at least that bound.

private theorem exists_mem_high_ge {t : IntervalTree} {x : Nat} (hx : x > 0) (h : realMaxHigh t ≥ x) : ∃ i ∈ keys t, i.high ≥ x := by induction t with | empty => simp [realMaxHigh, AugmentedTree.realAug, maxHighAug] at h omega | node left int _ right ihLeft ihRight => simp [realMaxHigh, AugmentedTree.realAug, maxHighAug] at h by_cases hi : int.high ≥ x · use int; simp [hi, keys_node] · have : realMaxHigh left ≥ x ∨ realMaxHigh right ≥ x := by simp [realMaxHigh, maxHighAug] at h ⊢ omega rcases this with h' | h' · obtain ⟨k, hk, hkHigh⟩ := ihLeft h' use k; simp [hk, hkHigh, keys_node] · obtain ⟨k, hk, hkHigh⟩ := ihRight h' use k; simp [hk, hkHigh, keys_node]

If the left subtree is non-empty, its max-high is at least the query low, and the current interval does not overlap, then any overlap in the right subtree forces an overlap in the left subtree. This is the key pruning invariant for interval search.

theorem overlap_left_of_right_overlap {left right : IntervalTree} {int q : Interval} (mx : Nat) (hB : IsBST (AugmentedTree.node left int mx right)) (hmax : realMaxHigh left ≥ q.low) (hcur : ¬ Interval.overlaps int q) (hright : ∃ j ∈ keys right, Interval.overlaps j q) : ∃ i ∈ keys left, Interval.overlaps i q := by rcases hright with ⟨j, hj, hov⟩ rw [Interval.overlaps_iff] at hov rcases hB with ⟨_hBL, _hBR, hLeftLE, hRightGE⟩ have h1 : int.low ≤ j.low := hRightGE j hj have h2 : j.low ≤ q.high := hov.1 have h3 : int.low ≤ q.high := by linarith have h4 : int.high < q.low := by by_contra h' have hov : Interval.overlaps int q = true := by rw [Interval.overlaps_iff] exact ⟨h3, by omega⟩ exact hcur hov have h5 : ∃ i ∈ keys left, i.high ≥ q.low := by have hx : q.low > 0 := by omega exact exists_mem_high_ge hx hmax rcases h5 with ⟨i, hi, hiHigh⟩ have h6 : i.low ≤ int.low := hLeftLE i hi have h7 : i.low ≤ q.high := by linarith use i, hi rw [Interval.overlaps_iff] exact ⟨h7, hiHigh⟩
end IntervalTreenamespace IntervalTree

hasOverlap distributes over a node in the obvious way.

@[simp] theorem hasOverlap_node {left right : IntervalTree} {int : Interval} {mx : Nat} {q : Interval} : hasOverlap (AugmentedTree.node left int mx right) q ↔ Interval.overlaps int q = true ∨ hasOverlap left q ∨ hasOverlap right q := by simp [hasOverlap, keys_node, Interval.overlaps_iff] constructor · rintro ⟨i, (hi | rfl | hi), hlow, hhigh⟩ · right; left; use i · left; exact ⟨hlow, hhigh⟩ · right; right; use i · rintro (⟨hlow, hhigh⟩ | ⟨i, hi, hlow, hhigh⟩ | ⟨i, hi, hlow, hhigh⟩) · use int; simp; exact ⟨hlow, hhigh⟩ · use i; simp [hi]; exact ⟨hlow, hhigh⟩ · use i; simp [hi]; exact ⟨hlow, hhigh⟩

The executable interval search returns only intervals that are in the tree and overlap the query.

theorem intervalSearch?_some_overlap {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) (i : Interval) : intervalSearch? t q = some i → i ∈ keys t ∧ Interval.overlaps i q := by induction t with | empty => simp [intervalSearch?] | node left int _ right ihLeft ihRight => intro h rcases hB with ⟨hBL, hBR, _hLeftLE, _hRightGE⟩ rcases hW with ⟨hWL, hWR, _hMax⟩ unfold intervalSearch? at h by_cases hO : Interval.overlaps int q = true · -- current interval overlaps simp [hO] at h cases h with | refl => constructor · simp [keys_node] · simp [hO] · -- current interval does not overlap by_cases hL : goLeft left q = true · -- go left simp [hO, hL] at h have hIh := ihLeft hBL hWL h constructor · simp [keys_node, hIh.1] · exact hIh.2 · -- go right simp [hO, hL] at h have hIh := ihRight hBR hWR h constructor · simp [keys_node, hIh.1] · exact hIh.2

If the executable interval search returns none, no interval in the tree overlaps the query.

theorem intervalSearch?_none_noOverlap {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) : intervalSearch? t q = none → ¬ hasOverlap t q := by induction t with | empty => simp [intervalSearch?, hasOverlap] | node left int mx right ihLeft ihRight => intro h rcases hB with ⟨hBL, hBR, hLeftLE, hRightGE⟩ rcases hW with ⟨hWL, hWR, _hMax⟩ unfold intervalSearch? at h by_cases hO : Interval.overlaps int q = true · -- current overlaps, but returned none: impossible simp [hO] at h · -- current interval does not overlap by_cases hL : goLeft left q = true · -- went left and got none simp [hO, hL] at h have hNoLeft : ¬ hasOverlap left q := ihLeft hBL hWL h intro hOv rcases hOv with ⟨j, hj, hov⟩ simp [keys_node] at hj rcases hj with (hjLeft | hjEq | hjRight) · -- overlap in left subtree exact hNoLeft ⟨j, hjLeft, hov⟩ · -- overlap with current interval rw [hjEq] at hov exact hO hov · -- overlap in right subtree forces one in left subtree have hRightEx : ∃ j ∈ keys right, Interval.overlaps j q := ⟨j, hjRight, hov⟩ have hL' := goLeft_true hL have hLeftOverlap := overlap_left_of_right_overlap mx ⟨hBL, hBR, hLeftLE, hRightGE⟩ (by rw [← storedMaxHigh_eq_realMaxHigh_of_wellAugmented hWL]; exact hL'.2) hO hRightEx rcases hLeftOverlap with ⟨k, hk, hkov⟩ exact hNoLeft ⟨k, hk, hkov⟩ · -- went right and got none have hL_false : goLeft left q = false := by simp [hL] simp [hO, hL] at h have hNoRight : ¬ hasOverlap right q := ihRight hBR hWR h intro hOv rcases hOv with ⟨j, hj, hov⟩ simp [keys_node] at hj rcases hj with (hjLeft | hjEq | hjRight) · -- no overlap possible in left subtree have hNoLeft : ¬ hasOverlap left q := by rcases goLeft_false hL_false with hEmpty | hmax · simp [hEmpty, hasOverlap] · have hmax' : realMaxHigh left < q.low := by rw [← storedMaxHigh_eq_realMaxHigh_of_wellAugmented hWL] exact hmax intro hOv' exact noOverlap_of_realMaxHigh_lt hmax' hOv' exact hNoLeft ⟨j, hjLeft, hov⟩ · -- overlap with current interval rw [hjEq] at hov exact hO hov · -- overlap in right subtree exact hNoRight ⟨j, hjRight, hov⟩

Combined correctness specification for interval search.

theorem intervalSearch?_spec {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) : (intervalSearch? t q = none ↔ ¬ hasOverlap t q) ∧ (∀ i, intervalSearch? t q = some i → i ∈ keys t ∧ Interval.overlaps i q) := by constructor · constructor · exact intervalSearch?_none_noOverlap hB hW q · intro hNoOverlap by_contra h have : intervalSearch? t q ≠ none := by simp [h] rcases Option.ne_none_iff_exists'.mp this with ⟨i, hi⟩ have := intervalSearch?_some_overlap hB hW q i hi exact hNoOverlap ⟨i, this.1, this.2⟩ · exact intervalSearch?_some_overlap hB hW q
end IntervalTree
Red-black bridge: maintaining an augmentation through Chapter 13 rotations

CLRS Theorem 14.1 in the red-black setting: a locally-computable augmentation can be maintained through the structural primitives used by red-black insertion and deletion. Chapter 13's red-black rotations and root recoloring are exactly those primitives. We show that any rotation-invariant augmentation's value is preserved by these operations — so a locally-recomputed cached field stays correct — reusing the same IsRotationInvariant law as the generic framework.

Note: red-black rotations are shape-restoring, not shape-preserving: they are applied mid-fixup where RedBlackShape is temporarily broken, so we do not (and cannot) claim a single rotation preserves RedBlackShape. RedBlackShape maintenance across a full insertion is Chapter 13's RBTree.redBlackShape_insert; the new content here is that the augmentation rides along invariantly under the same rotations and recoloring.

namespace RBBridgeopen CLRS.Chapter13

Inorder key list of a Chapter 13 red-black tree.

def rbKeys : RBTree → List Nat | .empty => [] | .node _ l k r => rbKeys l ++ [k] ++ rbKeys r

Semantic value of an augmentation on a red-black tree, computed from each key and its children (independent of node colors).

def rbRealAug {β : Type} [Inhabited β] (aug : Augmentation Nat β) : RBTree → β | .empty => aug.base | .node _ l k r => aug.combine k (rbRealAug aug l) (rbRealAug aug r)

Left rotation preserves the inorder key list.

theorem rbKeys_rotateLeft (t : RBTree) : rbKeys (RBTree.rotateLeft t) = rbKeys t := by cases t with | empty => rfl | node color a x right => cases right with | empty => rfl | node rc b y c => simp [RBTree.rotateLeft, rbKeys, List.append_assoc]

Right rotation preserves the inorder key list.

theorem rbKeys_rotateRight (t : RBTree) : rbKeys (RBTree.rotateRight t) = rbKeys t := by cases t with | empty => rfl | node color left y c => cases left with | empty => rfl | node lc a x b => simp [RBTree.rotateRight, rbKeys, List.append_assoc]

Left rotation preserves any rotation-invariant augmentation's value.

theorem rbRealAug_rotateLeft {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] (t : RBTree) : rbRealAug aug (RBTree.rotateLeft t) = rbRealAug aug t := by cases t with | empty => rfl | node color a x right => cases right with | empty => rfl | node rc b y c => simp only [RBTree.rotateLeft, rbRealAug] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves any rotation-invariant augmentation's value.

theorem rbRealAug_rotateRight {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] (t : RBTree) : rbRealAug aug (RBTree.rotateRight t) = rbRealAug aug t := by cases t with | empty => rfl | node color left y c => cases left with | empty => rfl | node lc a x b => simp only [RBTree.rotateRight, rbRealAug] rw [IsRotationInvariant.combine_rotate]

Root recoloring preserves the augmentation value.

theorem rbRealAug_repaintRoot {β : Type} [Inhabited β] (aug : Augmentation Nat β) (c : Color) (t : RBTree) : rbRealAug aug (RBTree.repaintRoot c t) = rbRealAug aug t := by cases t <;> simp [RBTree.repaintRoot, rbRealAug]

Red-black bridge (CLRS Theorem 14.1, red-black primitives). Every structural primitive used by red-black insertion and deletion — left and right rotation and root recoloring — preserves both the inorder key sequence and any rotation-invariant augmentation's value. Hence the augmentation can be maintained through the red-black operations by local recomputation, exactly as in the generic framework.

theorem rb_augmentation_bridge {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] : (∀ t, rbKeys (RBTree.rotateLeft t) = rbKeys t) ∧ (∀ t, rbKeys (RBTree.rotateRight t) = rbKeys t) ∧ (∀ t, rbRealAug aug (RBTree.rotateLeft t) = rbRealAug aug t) ∧ (∀ t, rbRealAug aug (RBTree.rotateRight t) = rbRealAug aug t) ∧ (∀ (c : Color) (t), rbRealAug aug (RBTree.repaintRoot c t) = rbRealAug aug t) := ⟨rbKeys_rotateLeft, rbKeys_rotateRight, rbRealAug_rotateLeft aug, rbRealAug_rotateRight aug, fun c t => rbRealAug_repaintRoot aug c t⟩

The size augmentation's value on a red-black tree is its node count.

theorem rbRealAug_sizeAug_eq_length (t : RBTree) : rbRealAug (sizeAug Nat) t = (rbKeys t).length := by induction t with | empty => rfl | node c l k r ihl ihr => simp only [rbRealAug, rbKeys] rw [ihl, ihr] simp only [sizeAug, List.length_append, List.length_cons, List.length_nil] omega
end RBBridge
General augmentation interface: an arbitrary augmentation through

executable red-black insertion

This section closes the "stored-field refinement" gap noted above, at the generic level. Section 14.1's OSRBTree threaded only the concrete subtree-size augmentation through Chapter 13's executable red-black insertion, in a bespoke, size-specific type. Here we thread an arbitrary Augmentation through the same Okasaki-style balancer, so both the order-statistic (size) and interval (max-high) augmentations are recovered as instances of a single generic interface.

The augmented red-black tree AugmentedRBTree caches, at every internal node, a node colour (reusing Chapter 13's CLRS.­Chapter13.­Color) and an augmentation value of type β. Every reconstructed node is built by the smart constructor AugmentedRBTree.mk, which recomputes the cached augmentation from its children via aug.combine. Two bridges connect this to the existing development:

  • AugmentedRBTree.wellAugmented_insert: the augmentation invariant survives balancing — inserting into a well-augmented tree yields a well-augmented tree, for any augmentation (CLRS 14.1 maintained through RB-INSERT).

  • AugmentedRBTree.toRB_insert: erasing the augmentation field commutes with insertion, so (for Nat keys) the augmented insert refines the executable Chapter 13 CLRS.­Chapter13.­RBTree.­insert exactly, transferring its shape and membership theorems.

The size and max-high fields are then recovered as instances via AugmentedRBTree.sizeAug_wellAugmented_insert and AugmentedRBTree.maxHighAug_wellAugmented_insert.

open CLRS.Chapter13 (Color RBTree)

A red-black tree augmented with a cached value of type β at every internal node. Each node stores a colour (reusing Chapter 13's CLRS.­Chapter13.­Color), a key of type α, a cached augmentation of type β, and two subtrees. This is the colour-carrying refinement of AugmentedTree, adding the field needed to run the Chapter 13 insertion balancer generically.

inductive AugmentedRBTree (α β : Type) where | empty : AugmentedRBTree α β | node : Color → AugmentedRBTree α β → α → β → AugmentedRBTree α β → AugmentedRBTree α β deriving Repr, DecidableEq
namespace AugmentedRBTreesection Genericvariable {α β : Type} [Inhabited β] (aug : Augmentation α β)

Inorder traversal of the keys, ignoring colours and cached augmentations.

def keys : AugmentedRBTree α β → List α | empty => [] | node _ l k _ r => keys l ++ [k] ++ keys r

The cached augmentation stored at the root; the empty tree uses aug.base.

def storedAug : AugmentedRBTree α β → β | empty => aug.base | node _ _ _ a _ => a

The mathematically correct augmentation, recomputed from the children.

def realAug : AugmentedRBTree α β → β | empty => aug.base | node _ l k _ r => aug.combine k (realAug l) (realAug r)

Every cached augmentation agrees with the recomputed one.

def WellAugmented : AugmentedRBTree α β → Prop | empty => True | node _ l k a r => WellAugmented l ∧ WellAugmented r ∧ a = aug.combine k (realAug aug l) (realAug aug r)

Smart constructor that recomputes the cached augmentation from the children. Every node produced by the red-black operations below is built with mk, which is why the WellAugmented invariant is preserved automatically.

def mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := node c l k (aug.combine k (realAug aug l) (realAug aug r)) r

mk recomputes the augmentation correctly.

theorem realAug_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : realAug aug (mk aug c l k r) = aug.combine k (realAug aug l) (realAug aug r) := rfl

mk preserves the inorder key sequence.

theorem keys_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : keys (mk aug c l k r) = keys l ++ [k] ++ keys r := rfl

The cached root augmentation of a mk node is the recomputed value.

theorem storedAug_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : storedAug aug (mk aug c l k r) = aug.combine k (realAug aug l) (realAug aug r) := rfl

A mk node is well-augmented whenever both children are.

theorem wellAugmented_mk {c : Color} {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (mk aug c l k r) := ⟨hl, hr, rfl⟩

A well-augmented tree has a correct root augmentation.

theorem storedAug_eq_realAug_of_wellAugmented {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug t = realAug aug t := by cases t with | empty => rfl | node c l k a r => exact h.2.2
Executable red-black operations with augmentation recomputation

Repaint the root black, keeping the cached augmentation fields.

def repaintBlack : AugmentedRBTree α β → AugmentedRBTree α β | empty => empty | node _ l k a r => node Color.black l k a r

Okasaki-style rebalance after insertion on the left child, recomputing augmentations. Mirrors CLRS.­Chapter13.­RBTree.­balanceLeft.

def balanceLeft (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match l with | node Color.red (node Color.red a w _ b) x _ c => mk aug Color.red (mk aug Color.black a w b) x (mk aug Color.black c y r) | node Color.red a w _ (node Color.red b x _ c) => mk aug Color.red (mk aug Color.black a w b) x (mk aug Color.black c y r) | _ => mk aug Color.black l y r

Okasaki-style rebalance after insertion on the right child, recomputing augmentations. Mirrors CLRS.­Chapter13.­RBTree.­balanceRight.

def balanceRight (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | node Color.red (node Color.red b x _ c) y' _ d => mk aug Color.red (mk aug Color.black l y b) x (mk aug Color.black c y' d) | node Color.red b x _ (node Color.red c y' _ d) => mk aug Color.red (mk aug Color.black l y b) x (mk aug Color.black c y' d) | _ => mk aug Color.black l y r

Insertion fixup: recurse down by the Boolean comparison lt, rebuilding and rebalancing with augmentation recomputation on the way back up. Mirrors CLRS.­Chapter13.­RBTree.­insertFixup.

def insertFixup (lt : α → α → Bool) (x : α) : AugmentedRBTree α β → AugmentedRBTree α β | empty => mk aug Color.red empty x empty | node c l y a r => if lt x y then if c = Color.black then balanceLeft aug (insertFixup lt x l) y r else mk aug Color.red (insertFixup lt x l) y r else if lt y x then if c = Color.black then balanceRight aug l y (insertFixup lt x r) else mk aug Color.red l y (insertFixup lt x r) else node c l y a r

Insert a key and repaint the root black.

def insert (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) : AugmentedRBTree α β := repaintBlack (insertFixup aug lt x t)
The augmentation invariant survives balancing (CLRS 14.1 through RB-INSERT)

Repainting the root black preserves the WellAugmented invariant.

theorem wellAugmented_repaintBlack {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (repaintBlack t) := by cases t with | empty => trivial | node c l k a r => exact ⟨h.1, h.2.1, h.2.2⟩

balanceLeft preserves the WellAugmented invariant.

theorem wellAugmented_balanceLeft {l : AugmentedRBTree α β} {y : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (balanceLeft aug l y r) := by unfold balanceLeft split · obtain ⟨⟨ha, hb, _⟩, hc, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) (wellAugmented_mk aug hc hr) · obtain ⟨ha, ⟨hb, hc, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) (wellAugmented_mk aug hc hr) · exact wellAugmented_mk aug hl hr

balanceRight preserves the WellAugmented invariant.

theorem wellAugmented_balanceRight {l : AugmentedRBTree α β} {y : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (balanceRight aug l y r) := by unfold balanceRight split · obtain ⟨⟨hb, hc, _⟩, hd, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hb) (wellAugmented_mk aug hc hd) · obtain ⟨hb, ⟨hc, hd, _⟩, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hb) (wellAugmented_mk aug hc hd) · exact wellAugmented_mk aug hl hr

insertFixup preserves the WellAugmented invariant.

theorem wellAugmented_insertFixup (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (insertFixup aug lt x t) := by induction t with | empty => simp only [insertFixup] exact wellAugmented_mk aug (by trivial) (by trivial) | node c l y a r ihl ihr => have hl : WellAugmented aug l := h.1 have hr : WellAugmented aug r := h.2.1 simp only [insertFixup] split · split · exact wellAugmented_balanceLeft aug (ihl hl) hr · exact wellAugmented_mk aug (ihl hl) hr · split · split · exact wellAugmented_balanceRight aug hl (ihr hr) · exact wellAugmented_mk aug hl (ihr hr) · exact h

Augmentation invariant through executable insertion (CLRS 14.1 through RB-INSERT). Inserting a key into a well-augmented augmented red-black tree produces a well-augmented tree: every cached augmentation field remains correct after the red-black rebalancing, for any Augmentation. This generalizes OSRBTree.wellSized_insert from the size field to an arbitrary augmentation.

theorem wellAugmented_insert (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (insert aug lt x t) := by unfold insert exact wellAugmented_repaintBlack aug (wellAugmented_insertFixup aug lt x h)

After insertion the cached root augmentation equals the recomputed value.

theorem storedAug_insert (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug (insert aug lt x t) = realAug aug (insert aug lt x t) := storedAug_eq_realAug_of_wellAugmented aug (wellAugmented_insert aug lt x h)
Executable red-black deletion with augmentation recomputation

The deletion pipeline mirrors OSRBTree but is generic in α, β, and aug. Every function recomputes augmentations via mk aug.

Repaint with arbitrary color, keeping cached augmentation fields.

def repaintRoot (c : Color) (t : AugmentedRBTree α β) : AugmentedRBTree α β := match t with | empty => empty | node _ l k a r => node c l k a r

Boolean black-root test.

def rootBlack : AugmentedRBTree α β → Bool | empty => true | node c _ _ _ _ => c = Color.black

Deletion re-balancer for a black-deficient left child.

def baldL (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match l with | node Color.red a x _ b => mk aug Color.red (mk aug Color.black a x b) k r | _ => match r with | node Color.black c y _ d => balanceRight aug l k (mk aug Color.red c y d) | node Color.red (node Color.black c y _ d) z _ e => mk aug Color.red (mk aug Color.black l k c) y (balanceRight aug d z (repaintRoot Color.red e)) | _ => mk aug Color.red l k r

Deletion re-balancer for a black-deficient right child.

def baldR (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | node Color.red c y _ d => mk aug Color.red l k (mk aug Color.black c y d) | _ => match l with | node Color.black a x _ b => balanceLeft aug (mk aug Color.red a x b) k r | node Color.red a x _ (node Color.black c y _ d) => mk aug Color.red (balanceLeft aug (repaintRoot Color.red a) x c) y (mk aug Color.black d k r) | _ => mk aug Color.red l k r

Find and remove the minimum key, recomputing augmentations on the way up.

def splitMin [Inhabited α] : AugmentedRBTree α β → α × AugmentedRBTree α β | empty => (default, empty) | node _ empty k _ r => (k, r) | node _ l k _ r => let (m, l') := splitMin l if rootBlack l then (m, baldL aug l' k r) else (m, mk aug Color.red l' k r)

Merge two trees, used when deleting a node with two children.

def join [Inhabited α] (l r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | empty => l | _ => match l with | empty => r | _ => let (m, r') := splitMin aug r if rootBlack r then baldR aug l m r' else mk aug Color.red l m r'

Recursive deletion.

def del [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) : AugmentedRBTree α β → AugmentedRBTree α β | empty => empty | node _c l y _a r => if lt x y then if rootBlack l then baldL aug (del x lt l) y r else mk aug Color.red (del x lt l) y r else if lt y x then if rootBlack r then baldR aug l y (del x lt r) else mk aug Color.red l y (del x lt r) else join aug l r

Delete a key and repaint the root black.

def delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) (t : AugmentedRBTree α β) : AugmentedRBTree α β := repaintBlack (del aug x lt t)
The augmentation invariant survives deletion
theorem wellAugmented_repaintRoot (c : Color) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (repaintRoot c t) := by cases t with | empty => trivial | node c' l k a r => exact ⟨h.1, h.2.1, h.2.2⟩theorem wellAugmented_baldL {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (baldL aug l k r) := by unfold baldL -- Use case analysis that matches the definition's pattern-matching structure cases l with | empty => cases r with | empty => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_balanceRight aug hl (wellAugmented_mk aug hc hd) | red => cases rl with | empty => exact wellAugmented_mk aug hl hr | node rlc rll rlk rla rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hc) (wellAugmented_balanceRight aug hd (wellAugmented_repaintRoot aug Color.red he)) | red => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | red => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) hr | black => cases r with | empty => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_balanceRight aug hl (wellAugmented_mk aug hc hd) | red => cases rl with | empty => exact wellAugmented_mk aug hl hr | node rlc rll rlk rla rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hc) (wellAugmented_balanceRight aug hd (wellAugmented_repaintRoot aug Color.red he)) | red => exact wellAugmented_mk aug hl hrtheorem wellAugmented_baldR {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (baldR aug l k r) := by unfold baldR cases r with | empty => cases l with | empty => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_balanceLeft aug (wellAugmented_mk aug ha hb) hr | red => cases lr with | empty => exact wellAugmented_mk aug hl hr | node lrc lrl lrk lra lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_balanceLeft aug (wellAugmented_repaintRoot aug Color.red ha) hc) (wellAugmented_mk aug hd hr) | red => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | red => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_mk aug hl (wellAugmented_mk aug hc hd) | black => cases l with | empty => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_balanceLeft aug (wellAugmented_mk aug ha hb) hr | red => cases lr with | empty => exact wellAugmented_mk aug hl hr | node lrc lrl lrk lra lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_balanceLeft aug (wellAugmented_repaintRoot aug Color.red ha) hc) (wellAugmented_mk aug hd hr) | red => exact wellAugmented_mk aug hl hr theorem wellAugmented_splitMin [Inhabited α] {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (splitMin aug t).2 := by induction t with | empty => trivial | node c l k a r ihl => cases l with | empty => exact h.2.1 | node lc ll lk la lr => have hws : WellAugmented aug (splitMin aug (node lc ll lk la lr)).2 := ihl h.1 by_cases hrb : rootBlack (node lc ll lk la lr) = true · have hsp : (splitMin aug (node c (node lc ll lk la lr) k a r)).2 = baldL aug (splitMin aug (node lc ll lk la lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellAugmented_baldL aug hws h.2.1 · have hsp : (splitMin aug (node c (node lc ll lk la lr) k a r)).2 = mk aug Color.red (splitMin aug (node lc ll lk la lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellAugmented_mk aug hws h.2.1theorem wellAugmented_join [Inhabited α] {l r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (join aug l r) := by unfold join split · exact hl · rename_i hneR split · exact hr · rename_i hneL dsimp by_cases hrb : rootBlack r = true · simp [hrb] exact wellAugmented_baldR aug hl (wellAugmented_splitMin aug hr) · simp [hrb] exact wellAugmented_mk aug hl (wellAugmented_splitMin aug hr)theorem wellAugmented_del [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (del aug x lt t) := by induction t with | empty => exact h | node c l y a r ihl ihr => have hl : WellAugmented aug l := h.1 have hr : WellAugmented aug r := h.2.1 simp only [del] split · split · exact wellAugmented_baldL aug (ihl hl) hr · exact wellAugmented_mk aug (ihl hl) hr · split · split · exact wellAugmented_baldR aug hl (ihr hr) · exact wellAugmented_mk aug hl (ihr hr) · exact wellAugmented_join aug hl hrtheorem wellAugmented_delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (delete aug x lt t) := by unfold delete exact wellAugmented_repaintBlack aug (wellAugmented_del aug x lt h)theorem storedAug_delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug (delete aug x lt t) = realAug aug (delete aug x lt t) := storedAug_eq_realAug_of_wellAugmented aug (wellAugmented_delete aug x lt h)end Genericsection Refinementvariable {β : Type} [Inhabited β] (aug : Augmentation Nat β)

Erase the cached augmentation field, projecting a Nat-keyed augmented red-black tree onto the Chapter 13 red-black tree.

def toRB : AugmentedRBTree Nat β → RBTree | empty => RBTree.empty | node c l k _ r => RBTree.node c (toRB l) k (toRB r)

The Nat strict-less-than comparison as a Bool, used to instantiate the generic insertion so that it refines Chapter 13's CLRS.­Chapter13.­RBTree.­insert.

def natLt (a b : Nat) : Bool := decide (a < b)

natLt decides strict less-than.

theorem natLt_true_iff {a b : Nat} : (natLt a b = true) ↔ a < b := by simp [natLt]

Erasing the augmentation of a mk node forgets only the cached value.

theorem toRB_mk (c : Color) (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (mk aug c l k r) = RBTree.node c (toRB l) k (toRB r) := rfl

Erasing the augmentation commutes with repainting the root black.

omit [Inhabited β] intheorem toRB_repaintBlack (t : AugmentedRBTree Nat β) : toRB (repaintBlack t) = RBTree.repaintRoot Color.black (toRB t) := by cases t with | empty => rfl | node c l k a r => rfl

Erasing the augmentation commutes with balanceLeft.

theorem toRB_balanceLeft (l : AugmentedRBTree Nat β) (y : Nat) (r : AugmentedRBTree Nat β) : toRB (balanceLeft aug l y r) = RBTree.balanceLeft (toRB l) y (toRB r) := by cases l with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the augmentation commutes with balanceRight.

theorem toRB_balanceRight (l : AugmentedRBTree Nat β) (y : Nat) (r : AugmentedRBTree Nat β) : toRB (balanceRight aug l y r) = RBTree.balanceRight (toRB l) y (toRB r) := by cases r with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the augmentation commutes with insertFixup (at natLt).

theorem toRB_insertFixup (x : Nat) (t : AugmentedRBTree Nat β) : toRB (insertFixup aug natLt x t) = RBTree.insertFixup x (toRB t) := by induction t with | empty => rfl | node c l y a r ihl ihr => simp only [insertFixup, RBTree.insertFixup, toRB, natLt_true_iff, gt_iff_lt, apply_ite toRB, toRB_mk, toRB_balanceLeft, toRB_balanceRight, ihl, ihr]

Refinement. The augmented insertion refines the executable Chapter 13 red-black insertion: erasing the cached augmentation turns insert (at natLt) into CLRS.­Chapter13.­RBTree.­insert. This generalizes OSRBTree.toRB_insert from the size field to an arbitrary augmentation.

theorem toRB_insert (x : Nat) (t : AugmentedRBTree Nat β) : toRB (insert aug natLt x t) = RBTree.insert x (toRB t) := by unfold insert RBTree.insert rw [toRB_repaintBlack, toRB_insertFixup]

Erasure relates keys membership to Chapter 13 tree membership.

omit [Inhabited β] intheorem inTree_toRB (y : Nat) (t : AugmentedRBTree Nat β) : RBTree.InTree y (toRB t) ↔ y ∈ keys t := by induction t with | empty => simp [toRB, RBTree.InTree, keys] | node c l k a r ihl ihr => simp only [toRB, RBTree.InTree, keys, List.append_assoc, List.singleton_append, List.mem_append, List.mem_cons, ihl, ihr] tauto

Through the refinement, the Chapter 13 red-black shape invariant is maintained by the augmented insertion.

theorem redBlackShape_toRB_insert (x : Nat) {t : AugmentedRBTree Nat β} (h : RBTree.RedBlackShape (toRB t)) : RBTree.RedBlackShape (toRB (insert aug natLt x t)) := by rw [toRB_insert] exact RBTree.redBlackShape_insert h

Through the refinement, insertion preserves membership (as an inorder key).

theorem mem_keys_insert (x y : Nat) (t : AugmentedRBTree Nat β) : y ∈ keys (insert aug natLt x t) ↔ y = x ∨ y ∈ keys t := by simp only [← inTree_toRB, toRB_insert, RBTree.inTree_insert_iff]
Refinement of the deletion pipeline

Erasing the augmentation commutes with repainting the root in any colour.

omit [Inhabited β] intheorem toRB_repaintRoot (c : Color) (t : AugmentedRBTree Nat β) : toRB (repaintRoot c t) = RBTree.repaintRoot c (toRB t) := by cases t <;> rfl

Erasing the augmentation preserves the boolean black-root test.

omit [Inhabited β] intheorem rootBlack_toRB (t : AugmentedRBTree Nat β) : RBTree.rootBlack (toRB t) = rootBlack t := by cases t <;> rfl

Erasing the augmentation preserves the empty/non-empty distinction.

omit [Inhabited β] intheorem toRB_empty_iff (t : AugmentedRBTree Nat β) : toRB t = RBTree.empty ↔ t = AugmentedRBTree.empty := by cases t <;> simp [toRB]

Erasing the augmentation commutes with baldL.

theorem toRB_baldL (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (baldL aug l k r) = RBTree.baldL (toRB l) k (toRB r) := by cases l with | empty => cases r with | empty => rfl | node rc cl cy sc cr => cases rc with | black => simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases cl with | empty => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node clc cll clk cls clr => cases clc <;> simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot] | node lc a x s b => cases lc with | red => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | black => cases r with | empty => rfl | node rc cl cy sc cr => cases rc with | black => simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases cl with | empty => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node clc cll clk cls clr => cases clc <;> simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot]

Erasing the augmentation commutes with baldR.

theorem toRB_baldR (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (baldR aug l k r) = RBTree.baldR (toRB l) k (toRB r) := by cases r with | empty => cases l with | empty => rfl | node lc a x s b => cases lc with | black => simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases b with | empty => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node bc bl bk bs br => cases bc <;> simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot] | node rc c y s d => cases rc with | red => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | black => cases l with | empty => rfl | node lc a x s b => cases lc with | black => simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases b with | empty => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node bc bl bk bs br => cases bc <;> simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot]

Erasing the augmentation commutes with the minimum key of splitMin.

theorem toRB_splitMin_min (t : AugmentedRBTree Nat β) : (RBTree.splitMin (toRB t)).1 = (splitMin aug t).1 := by induction t with | empty => rfl | node c l k a r ihl => cases l with | empty => rfl | node lc ll lk la lr => cases lc with | black => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack] using ihl | red => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack] using ihl

Erasing the augmentation commutes with the tree component of splitMin.

theorem toRB_splitMin_tree (t : AugmentedRBTree Nat β) : toRB (splitMin aug t).2 = (RBTree.splitMin (toRB t)).2 := by induction t with | empty => rfl | node c l k a r ihl => cases l with | empty => rfl | node lc ll lk la lr => cases lc with | black => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack, toRB_baldL] using (congrArg (fun t : RBTree => RBTree.baldL t k r.toRB) ihl) | red => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack, toRB_mk] using (congrArg (fun t : RBTree => RBTree.node Color.red t k r.toRB) ihl)

Erasing the augmentation commutes with join.

theorem toRB_join (l r : AugmentedRBTree Nat β) : toRB (join aug l r) = RBTree.join (toRB l) (toRB r) := by cases r with | empty => cases l with | empty => rfl | node _ _ _ _ _ => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵toRB_splitMin_min, toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, This simp argument is unused: toRB_splitMin_min Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_min, This simp argument is unused: toRB_splitMin_tree Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵t̵r̵e̵e̵]̵t̲o̲R̲B̲_̲s̲p̲l̲i̲t̲M̲i̲n̲_̲m̲i̲n̲]̲ Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_tree] | node rc rl rk sr rr => cases l with | empty => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵toRB_splitMin_min, toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, This simp argument is unused: toRB_splitMin_min Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_min, This simp argument is unused: toRB_splitMin_tree Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵t̵r̵e̵e̵]̵t̲o̲R̲B̲_̲s̲p̲l̲i̲t̲M̲i̲n̲_̲m̲i̲n̲]̲ Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_tree] | node lc ll lk sl lr => cases rc with | black => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵rootBlack, rootBlack_toRB, ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲toRB_splitMin_tree, toRB_baldR] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, rootBlack, rootBlack_toRB, toRB_splitMin_tree, toRB_baldR] rw [toRB_splitMin_min] | red => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵rootBlack, rootBlack_toRB, ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲toRB_splitMin_tree, toRB_mk] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, rootBlack, rootBlack_toRB, toRB_splitMin_tree, toRB_mk] rw [toRB_splitMin_min]

Erasing the augmentation commutes with recursive deletion.

theorem toRB_del (x : Nat) (t : AugmentedRBTree Nat β) : toRB (del aug x natLt t) = RBTree.del x (toRB t) := by induction t with | empty => rfl | node c l y a r ihl ihr => by_cases hxy : x < y · have hxy' : natLt x y = true := natLt_true_iff.mpr hxy simp [del, RBTree.del, toRB, apply_ite toRB, rootBlack, hxy, hxy', rootBlack_toRB, toRB_baldL, toRB_mk, ihl] · by_cases hyx : y < x · have hxy' : natLt x y = false := by simp [natLt, hxy] have hyx' : natLt y x = true := natLt_true_iff.mpr hyx simp [del, RBTree.del, toRB, apply_ite toRB, rootBlack, hxy, hyx, hxy', hyx', rootBlack_toRB, toRB_baldR, toRB_mk, ihr] · have hxy' : natLt x y = false := by simp [natLt, hxy] have hyx' : natLt y x = false := by simp [natLt, hyx] simp [del, RBTree.del, toRB, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [del, RBTree.del, toRB, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵hxy, hyx, hxy', hyx', toRB_join] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, hxy, hyx, hxy', hyx', toRB_join]

Deletion refinement. The augmented deletion refines the executable Chapter 13 red-black deletion: erasing the cached augmentation turns delete (at natLt) into CLRS.­Chapter13.­RBTree.­delete. This extends OSRBTree.toRB_delete from the size field to an arbitrary augmentation, closing the deletion half of the generic AugmentedRBTree refinement.

theorem toRB_delete (x : Nat) (t : AugmentedRBTree Nat β) : toRB (delete aug x natLt t) = RBTree.delete x (toRB t) := by unfold delete RBTree.delete rw [toRB_repaintBlack, toRB_del]
end Refinementsection Instances
Instance 1: the order-statistic (subtree-size) augmentation

Taking aug := sizeAug Nat recovers the order-statistic tree of §14.1: the cached field is the subtree node count, the invariant survives the executable red-black insertion, and the augmentation-erasing projection refines Chapter 13's CLRS.­Chapter13.­RBTree.­insert. This makes OSRBTree a special case of the generic interface rather than a bespoke copy.

Order-statistic instance. The subtree-size augmentation is maintained through the generic executable red-black insertion (CLRS 14.1 for size).

theorem sizeAug_wellAugmented_insert (x : Nat) {t : AugmentedRBTree Nat Nat} (h : WellAugmented (sizeAug Nat) t) : WellAugmented (sizeAug Nat) (insert (sizeAug Nat) natLt x t) := wellAugmented_insert (sizeAug Nat) natLt x h

The size augmentation's recomputed value is exactly the node count.

theorem sizeAug_realAug_eq_length (t : AugmentedRBTree Nat Nat) : realAug (sizeAug Nat) t = (keys t).length := by induction t with | empty => rfl | node c l k a r ihl ihr => simp only [realAug] rw [ihl, ihr] simp only [sizeAug, keys, List.length_append, List.length_cons, List.length_nil] omega

Order-statistic refinement. Erasing the size field turns the generic size-augmented insertion into Chapter 13's CLRS.­Chapter13.­RBTree.­insert.

theorem sizeAug_toRB_insert (x : Nat) (t : AugmentedRBTree Nat Nat) : toRB (insert (sizeAug Nat) natLt x t) = RBTree.insert x (toRB t) := toRB_insert (sizeAug Nat) x t
Instance 2: the interval-tree (maximum-high-endpoint) augmentation

Taking aug := IntervalTree.maxHighAug recovers interval trees: the cached field is the subtree's maximum high endpoint, maintained through the same generic executable insertion, with distinct intervals ordered lexicographically by low and high endpoints.

Compare intervals lexicographically by low endpoint, then high endpoint. Distinct intervals with equal low endpoints remain distinct insertion keys.

def intervalLt (i j : Interval) : Bool := decide (i.low < j.low ∨ (i.low = j.low ∧ i.high < j.high))

Interval-tree instance. The maximum-high-endpoint augmentation is maintained through the generic executable red-black insertion.

end Instancesend AugmentedRBTree
Imports

17.3. Interval Trees

Interval insertion compares both endpoints lexicographically, retaining distinct intervals with equal low endpoints. The generic insertion companion proves complete-key membership and preservation of compatible transitive orderings. Its interval instance preserves strict lexicographic BST order and the weaker low-endpoint ordering used by the static search algorithm.

Color erasure preserves keys, max-high augmentation, and the static BST predicate. intervalSearch_insert_spec derives both post-insertion invariants from the input and proves complete search over the enlarged interval set: a failed search means neither the inserted interval nor any old interval overlaps the query. A successful search returns an old or newly inserted key that overlaps. intervalSearch_insert_finds guarantees success whenever the newly inserted interval overlaps the query.

The historical intervalSearch_after_update theorem remains the augmentation-only part of that bridge. The search-height analysis uses a conservative descent budget, which continues its accounting after a root match; its logarithmic bound assumes the supplied red-black shape invariant. These search results do not constitute a new red-black shape-preservation theorem for the interval insertion pipeline.

namespace CLRSnamespace Chapter14namespace IntervalTree

The height of an interval tree (maximum depth of the augmented tree).

def intervalHeight : IntervalTree → Nat | AugmentedTree.empty => 0 | AugmentedTree.node l _ _ r => 1 + max (intervalHeight l) (intervalHeight r)

Conservative descent budget for intervalSearch?: one node per level, continuing the accounting even when a matching interval stops search.

def intervalSearchCost : IntervalTree → Interval → Nat | AugmentedTree.empty, _ => 0 | AugmentedTree.node l _ _ r, q => 1 + if goLeft l q then intervalSearchCost l q else intervalSearchCost r q

The conservative interval-search descent budget is bounded by height plus one.

theorem intervalSearchCost_le_height (t : IntervalTree) (q : Interval) : intervalSearchCost t q ≤ intervalHeight t + 1 := by induction t generalizing q with | empty => simp [intervalSearchCost, intervalHeight] | node l int a r ihl ihr => simp only [intervalSearchCost, intervalHeight] by_cases h : goLeft l q · simp [h] have ih := ihl q have hmax : intervalHeight l ≤ max (intervalHeight l) (intervalHeight r) := Nat.le_max_left _ _ omega · simp [h] have ih := ihr q have hmax : intervalHeight r ≤ max (intervalHeight l) (intervalHeight r) := Nat.le_max_right _ _ omega
end IntervalTreenamespace AugmentedRBTree

Erase the colors of a dynamic augmented red-black interval tree, projecting it onto the static IntervalTree.

Erasing colors preserves the mathematical max-high augmentation.

theorem realAug_toIntervalTree (t : AugmentedRBTree Interval Nat) : AugmentedTree.realAug IntervalTree.maxHighAug (toIntervalTree t) = realAug IntervalTree.maxHighAug t := by induction t with | empty => rfl | node c l k a r ihl ihr => simp [toIntervalTree, AugmentedTree.realAug, realAug, ihl, ihr]

The erasure of a well-augmented dynamic interval tree is a well-augmented static interval tree (with the max-high augmentation).

theorem wellAugmented_toIntervalTree {t : AugmentedRBTree Interval Nat} (h : WellAugmented IntervalTree.maxHighAug t) : IntervalTree.WellAugmented (toIntervalTree t) := by induction t with | empty => simp [toIntervalTree, IntervalTree.WellAugmented] | node c l k a r ihl ihr => obtain ⟨hL, hR, ha⟩ := h change AugmentedTree.WellAugmented IntervalTree.maxHighAug (AugmentedTree.node (toIntervalTree l) k a (toIntervalTree r)) constructor · exact ihl hL constructor · exact ihr hR · simp only [AugmentedTree.realAug] rw [realAug_toIntervalTree l, realAug_toIntervalTree r] exact ha
end AugmentedRBTree

The dynamic/static bridge and search-after-update

The augmentation part of the insertion bridge. The stronger intervalSearch_insert_spec below also derives BST preservation and search correctness from the input invariants.

The Interval-keyed O(log n) search bound

open CLRS.Chapter13 (RBTree)namespace AugmentedRBTree

Erase the interval keys (keeping their low endpoint) and the cached augmentation, projecting an Interval-keyed augmented red-black tree onto the Chapter 13 RBTree.

def toRB_low : AugmentedRBTree Interval Nat → RBTree | empty => RBTree.empty | node c l k _ r => RBTree.node c (toRB_low l) k.low (toRB_low r)

The height of the static interval erasure equals the height of the low-keyed red-black erasure: heights depend on neither keys, colors, nor the cached augmentation.

theorem intervalHeight_eq_toRB_height (t : AugmentedRBTree Interval Nat) : IntervalTree.intervalHeight (toIntervalTree t) = RBTree.height (toRB_low t) := by induction t with | empty => rfl | node c l k a r ihl ihr => simp [toIntervalTree, toRB_low, IntervalTree.intervalHeight, RBTree.height, ihl, ihr]
end AugmentedRBTree

Interval search runs in O(log n). On an Interval-keyed augmented red-black tree with n nodes, the conservative search descent budget is at most 2 log₂(n+1) + 1 node visits, composing intervalSearchCost_le_height with the red-black height bound (RBTree.height_log_bound) via AugmentedRBTree.intervalHeight_eq_toRB_height.

namespace AugmentedRBTree

Color erasure preserves every complete interval key.

theorem keys_toIntervalTree (t : AugmentedRBTree Interval Nat) : IntervalTree.keys (toIntervalTree t) = keys t := by induction t with | empty => rfl | node c l k a r ihl ihr => simp only [toIntervalTree, IntervalTree.keys_node, keys, ihl, ihr]

The static search BST predicate is exactly weak inorder low ordering.

theorem isBST_toIntervalTree_iff (t : AugmentedRBTree Interval Nat) : IntervalTree.IsBST (toIntervalTree t) ↔ LowOrdered t := by induction t with | empty => simp [toIntervalTree, IntervalTree.IsBST, LowOrdered, Ordered, keys] | node c l k a r ihl ihr => rw [show LowOrdered (.node c l k a r) ↔ LowOrdered l ∧ LowOrdered r ∧ (∀ i ∈ keys l, i.low ≤ k.low) ∧ (∀ i ∈ keys r, k.low ≤ i.low) from ordered_node (fun i j : Interval => i.low ≤ j.low) (fun _ _ _ => le_trans) c l k a r] simp only [toIntervalTree, IntervalTree.IsBST, ihl, ihr, IntervalTree.allLowLE, IntervalTree.allLowGE, keys_toIntervalTree]

Insertion preserves the precise BST premise needed by static interval search.

Strict lexicographic interval BSTs satisfy the static search ordering.

theorem IntervalBST.toIntervalTree {t : AugmentedRBTree Interval Nat} (h : IntervalBST t) : IntervalTree.IsBST (toIntervalTree t) := (isBST_toIntervalTree_iff t).mpr h.lowOrdered

Overlaps after insertion are exactly the new interval's overlaps plus old ones.

theorem hasOverlap_interval_insert (q query : Interval) (t : AugmentedRBTree Interval Nat) : IntervalTree.hasOverlap (toIntervalTree (insert IntervalTree.maxHighAug intervalLt q t)) query ↔ Interval.overlaps q query = true ∨ IntervalTree.hasOverlap (toIntervalTree t) query := by simp only [IntervalTree.hasOverlap, keys_toIntervalTree, mem_keys_interval_insert] constructor · rintro ⟨i, hi, hov⟩ rcases hi with rfl | hi · exact Or.inl hov · exact Or.inr ⟨i, hi, hov⟩ · rintro (hq | ⟨i, hi, hov⟩) · exact ⟨q, Or.inl rfl, hq⟩ · exact ⟨i, Or.inr hi, hov⟩
end AugmentedRBTree

Search after actual insertion is complete for the enlarged interval set and returns only a stored overlapping interval. Both input invariants are explicit, and their postconditions are derived from the insertion execution.

An inserted interval overlapping the query guarantees a successful search.

theorem intervalSearch_insert_finds (q query : Interval) {t : AugmentedRBTree Interval Nat} (hB : IntervalTree.IsBST (AugmentedRBTree.toIntervalTree t)) (hW : AugmentedRBTree.WellAugmented IntervalTree.maxHighAug t) (hq : Interval.overlaps q query = true) : ∃ i, IntervalTree.intervalSearch? (AugmentedRBTree.toIntervalTree (AugmentedRBTree.insert IntervalTree.maxHighAug AugmentedRBTree.intervalLt q t)) query = some i ∧ (i = q ∨ i ∈ AugmentedRBTree.keys t) ∧ Interval.overlaps i query = true := by have hs := intervalSearch_insert_spec q query hB hW have hn : IntervalTree.intervalSearch? (AugmentedRBTree.toIntervalTree (AugmentedRBTree.insert IntervalTree.maxHighAug AugmentedRBTree.intervalLt q t)) query ≠ none := by intro he exact hs.1.mp he (Or.inl hq) obtain ⟨i, hi⟩ := Option.ne_none_iff_exists'.mp hn exact ⟨i, hi, hs.2 i hi⟩
end Chapter14end CLRS

Definitions and proofs

CLRSLean.FourthEdition.Chapter_17.Section_17_3_Interval_Trees.CachedInsertion

Interval search after the counted cached-field insertion

namespace CLRS.Chapter14

The counted cached-field insertion inherits complete post-insert search semantics.

theorem intervalSearch_cachedInsert_spec (q query : Interval) {t : AugmentedRBTree Interval Nat} (hB : IntervalTree.IsBST (AugmentedRBTree.toIntervalTree t)) (hW : AugmentedRBTree.WellAugmented IntervalTree.maxHighAug t) : let run := AugmentationExecution.insert IntervalTree.maxHighAug AugmentedRBTree.intervalLt q t let updated := AugmentedRBTree.toIntervalTree run.tree (IntervalTree.intervalSearch? updated query = none ↔ ¬ (Interval.overlaps q query = true ∨ IntervalTree.hasOverlap (AugmentedRBTree.toIntervalTree t) query)) ∧ (∀ i, IntervalTree.intervalSearch? updated query = some i → (i = q ∨ i ∈ AugmentedRBTree.keys t) ∧ Interval.overlaps i query = true) := by dsimp only rw [AugmentationExecution.insert_refines _ _ _ _ hW] exact intervalSearch_insert_spec q query hB hW
end CLRS.Chapter14

CLRSLean.FourthEdition.Chapter_17.Section_17_3_Interval_Trees.Insertion

Complete-key interval insertion

Generic augmented red-black insertion preserves membership when comparator ties imply key equality. Balancing preserves the exact inorder key list, so insertion also preserves any compatible transitive ordering relation. The interval instance uses strict lexicographic order on both endpoints, and separately preserves the weak low-endpoint ordering required by the static search proof.

namespace CLRS.Chapter14.AugmentedRBTreeopen CLRS.Chapter13 (Color)variable {α β : Type} [Inhabited β] (aug : Augmentation α β)

Root repainting preserves the full inorder key list.

omit [Inhabited β] in@[simp] theorem keys_repaintBlack_generic (t : AugmentedRBTree α β) : keys (repaintBlack t) = keys t := by cases t <;> rfl

Left insertion balancing preserves the full inorder key list.

@[simp] theorem keys_balanceLeft_generic (l : AugmentedRBTree α β) (k : α) (r) : keys (balanceLeft aug l k r) = keys l ++ [k] ++ keys r := by unfold balanceLeft split <;> simp [mk, keys, List.append_assoc]

Right insertion balancing preserves the full inorder key list.

@[simp] theorem keys_balanceRight_generic (l : AugmentedRBTree α β) (k : α) (r) : keys (balanceRight aug l k r) = keys l ++ [k] ++ keys r := by unfold balanceRight split <;> simp [mk, keys, List.append_assoc]

Fixup inserts the requested complete key and preserves all old keys.

theorem mem_keys_insertFixup_of_compare (lt : α → α → Bool) (hsep : ∀ x y, lt x y ≠ true → lt y x ≠ true → x = y) (x y : α) (t : AugmentedRBTree α β) : y ∈ keys (insertFixup aug lt x t) ↔ y = x ∨ y ∈ keys t := by induction t with | empty => simp [insertFixup, mk, keys] | node c l k a r ihl ihr => simp only [insertFixup] split · split <;> simp only [keys_balanceLeft_generic, keys_mk, List.mem_append, List.mem_singleton, ihl, keys] <;> tauto · rename_i hx split · split <;> simp only [keys_balanceRight_generic, keys_mk, List.mem_append, List.mem_singleton, ihr, keys] <;> tauto · rename_i hk have he := hsep x k hx hk subst x simp only [keys, List.mem_append, List.mem_singleton] tauto

Comparator ties must identify equal keys for set insertion membership.

theorem mem_keys_insert_of_compare (lt : α → α → Bool) (hsep : ∀ x y, lt x y ≠ true → lt y x ≠ true → x = y) (x y : α) (t : AugmentedRBTree α β) : y ∈ keys (insert aug lt x t) ↔ y = x ∨ y ∈ keys t := by simpa [insert] using mem_keys_insertFixup_of_compare aug lt hsep x y t
private theorem pairwise_middle (rel : α → α → Prop) (htrans : ∀ ⦃x y z⦄, rel x y → rel y z → rel x z) (L R : List α) (k : α) : (L ++ [k] ++ R).Pairwise (fun x y => rel x y) ↔ L.Pairwise (fun x y => rel x y) ∧ R.Pairwise (fun x y => rel x y) ∧ (∀ x ∈ L, rel x k) ∧ (∀ y ∈ R, rel k y) := by simp only [List.pairwise_append, List.pairwise_singleton, List.mem_append, List.mem_singleton] constructor · rintro ⟨⟨hL, _, hLk⟩, hR, hcross⟩ exact ⟨hL, hR, fun x hx => hLk x hx k rfl, fun y hy => hcross k (Or.inr rfl) y hy⟩ · rintro ⟨hL, hR, hLk, hkR⟩ refine ⟨⟨hL, trivial, ?_⟩, hR, ?_⟩ · rintro x hx y rfl exact hLk x hx · rintro x (hx | rfl) y hy · exact htrans (hLk x hx) (hkR y hy) · exact hkR y hy

All inorder keys satisfy a supplied ordering relation.

def Ordered (rel : α → α → Prop) (t : AugmentedRBTree α β) : Prop := (keys t).Pairwise (fun x y => rel x y)

A transitive inorder relation is equivalent to the recursive BST bounds.

omit [Inhabited β] intheorem ordered_node (rel : α → α → Prop) (htrans : ∀ ⦃x y z⦄, rel x y → rel y z → rel x z) (c l k a r) : Ordered rel (.node c l k a r : AugmentedRBTree α β) ↔ Ordered rel l ∧ Ordered rel r ∧ (∀ x ∈ keys l, rel x k) ∧ (∀ y ∈ keys r, rel k y) := pairwise_middle rel htrans (keys l) (keys r) k

Fixup preserves any transitive relation compatible with comparison.

theorem ordered_insertFixup (lt : α → α → Bool) (hsep : ∀ x y, lt x y ≠ true → lt y x ≠ true → x = y) (rel : α → α → Prop) (htrans : ∀ ⦃x y z⦄, rel x y → rel y z → rel x z) (hcompat : ∀ x y, lt x y = true → rel x y) (x : α) {t : AugmentedRBTree α β} (ht : Ordered rel t) : Ordered rel (insertFixup aug lt x t) := by induction t with | empty => simp [insertFixup, Ordered, mk, keys] | node c l k a r ihl ihr => obtain ⟨hl, hr, hLk, hkR⟩ := (ordered_node rel htrans c l k a r).mp ht simp only [insertFixup] split · rename_i hx have hnew : (keys (insertFixup aug lt x l) ++ [k] ++ keys r).Pairwise (fun u v => rel u v) := by apply (pairwise_middle rel htrans _ _ _).mpr refine ⟨ihl hl, hr, ?_, hkR⟩ intro y hy rcases (mem_keys_insertFixup_of_compare aug lt hsep x y l).mp hy with rfl | hy · exact hcompat y k hx · exact hLk y hy split <;> simpa only [Ordered, keys_balanceLeft_generic, keys_mk] using hnew · split · rename_i hx have hnew : (keys l ++ [k] ++ keys (insertFixup aug lt x r)).Pairwise (fun u v => rel u v) := by apply (pairwise_middle rel htrans _ _ _).mpr refine ⟨hl, ihr hr, hLk, ?_⟩ intro y hy rcases (mem_keys_insertFixup_of_compare aug lt hsep x y r).mp hy with rfl | hy · exact hcompat k y hx · exact hkR y hy split <;> simpa only [Ordered, keys_balanceRight_generic, keys_mk] using hnew · exact ht

Complete insertion preserves the supplied BST ordering.

theorem ordered_insert (lt : α → α → Bool) (hsep : ∀ x y, lt x y ≠ true → lt y x ≠ true → x = y) (rel : α → α → Prop) (htrans : ∀ ⦃x y z⦄, rel x y → rel y z → rel x z) (hcompat : ∀ x y, lt x y = true → rel x y) (x : α) {t : AugmentedRBTree α β} (ht : Ordered rel t) : Ordered rel (insert aug lt x t) := by simpa only [insert, Ordered, keys_repaintBlack_generic] using ordered_insertFixup aug lt hsep rel htrans hcompat x ht

Equal comparison keys are exactly equal intervals, including high endpoints.

theorem intervalLt_separates (i j : Interval) (hij : intervalLt i j ≠ true) (hji : intervalLt j i ≠ true) : i = j := by simp only [intervalLt, ne_eq, decide_eq_true_eq, Interval.low, Interval.high] at hij hji apply Prod.ext <;> omega

The interval comparator is the strict lexicographic order.

theorem intervalLt_trans : ∀ ⦃i j k⦄, intervalLt i j = true → intervalLt j k = true → intervalLt i k = true := by intro i j k hij hjk simp only [intervalLt, decide_eq_true_eq, Interval.low, Interval.high] at * omega

Lexicographic ordering implies the weak low-endpoint ordering used by search.

theorem intervalLt_low_le {i j : Interval} (h : intervalLt i j = true) : i.low ≤ j.low := by simp only [intervalLt, decide_eq_true_eq] at h omega

Strict lexicographic BST ordering of the complete interval keys.

def IntervalBST (t : AugmentedRBTree Interval Nat) : Prop := Ordered (fun i j => intervalLt i j = true) t

The weaker ordering needed by the static interval-search algorithm.

def LowOrdered (t : AugmentedRBTree Interval Nat) : Prop := Ordered (fun i j : Interval => i.low ≤ j.low) t

Interval insertion retains all old intervals and adds the complete new interval.

theorem mem_keys_interval_insert (q i : Interval) (t : AugmentedRBTree Interval Nat) : i ∈ keys (insert IntervalTree.maxHighAug intervalLt q t) ↔ i = q ∨ i ∈ keys t := mem_keys_insert_of_compare IntervalTree.maxHighAug intervalLt intervalLt_separates q i t

Insertion preserves strict lexicographic interval BST ordering.

Weak low ordering is preserved even for inputs not strictly lexicographically ordered.

theorem lowOrdered_insert (q : Interval) {t : AugmentedRBTree Interval Nat} (ht : LowOrdered t) : LowOrdered (insert IntervalTree.maxHighAug intervalLt q t) := ordered_insert IntervalTree.maxHighAug intervalLt intervalLt_separates (fun i j : Interval => i.low ≤ j.low) (fun _ _ _ h₁ h₂ => le_trans h₁ h₂) (fun _ _ h => intervalLt_low_le h) q ht

A lexicographic interval BST satisfies the ordering used by interval search.

theorem IntervalBST.lowOrdered {t : AugmentedRBTree Interval Nat} (ht : IntervalBST t) : LowOrdered t := List.Pairwise.imp (fun h => intervalLt_low_le h) ht
end CLRS.Chapter14.AugmentedRBTree

CLRSLean.Chapter_14.Section_14_3_Interval_Trees

Section 14.3 - Interval trees

This file formalizes the second classic augmentation from CLRS Chapter 14: interval trees. Each node stores an interval and a cached maximum high endpoint of its subtree. The search algorithm uses the cached maximum to prune the left subtree when no overlap is possible there.

To make the pattern reusable, we first define a small generic augmentation framework: an AugmentedTree α β carries a value of type α at each node and a cached augmentation of type β. An Augmentation α β provides the empty default and the local recombination function. The framework proves that recomputing the augmentation from the children preserves the WellAugmented invariant, and that rotations preserve the inorder key sequence and the WellAugmented invariant for augmentations whose combine operator behaves like max (which is the case for the interval-tree instantiation). We then instantiate it to interval trees with the max-high augmentation, and show that the executable intervalSearch? is correct on well-augmented BSTs.

Main results:

  • Generic AugmentedTree lemmas: keys_recompute, realAug_recompute, recompute_wellAugmented, rotateLeft_wellAugmented, rotateRight_wellAugmented.

  • Interval-tree correctness: intervalSearch?_some_overlap and intervalSearch?_none_noOverlap (combined as intervalSearch?_spec).

  • General augmentation theorem (CLRS Theorem 14.1): augmentation_theorem packages that rotations, recomputation, and generic BST insert maintain the WellAugmented invariant and the semantic augmentation for any rotation-invariant augmentation.

  • Size augmentation instance: sizeAug with realAug_sizeAug_eq_length, showing order-statistic size caching is an instance of the same framework.

  • Red-black bridge: rb_augmentation_bridge shows Chapter 13's red-black rotations and root recoloring preserve any rotation-invariant augmentation's value (and the inorder key list), so the augmentation is maintainable through the red-black operations.

  • General augmentation interface: AugmentedRBTree threads an arbitrary Augmentation through an executable red-black insertion; its smart constructor AugmentedRBTree.mk recomputes the cached value, so AugmentedRBTree.wellAugmented_insert shows the invariant survives balancing and AugmentedRBTree.toRB_insert shows the augmentation-erasing projection refines Chapter 13's executable RBTree.insert. Both the size and interval instances are recovered from it (AugmentedRBTree.sizeAug_wellAugmented_insert, AugmentedRBTree.maxHighAug_wellAugmented_insert).

Status: the static interval-search specification, generic local augmentation invariant, and arbitrary cached-field insertion/deletion pipeline are proved. The complete fourth-edition interval-tree and augmentation interfaces remain partial. Missing bridges include the constant-time-combine asymptotic theorem, combined BST/red-black/augmentation preservation, and interval-specific update semantics connecting AugmentedRBTree to the separate static IntervalTree search model. The current low-endpoint-only comparator also needs an equal-low policy before arbitrary intervals can be inserted distinctly.

namespace CLRSnamespace Chapter14
Generic augmented trees

An augmentation schema for a binary tree: a default for the empty tree and a local recombination function.

structure Augmentation (α β : Type) [Inhabited β] where base : β combine : α → β → β → β

Typeclass asserting that an augmentation's combine operation satisfies the rotation-invariance law required for BST rotations to preserve the cached augmentation. The max-high augmentation is the motivating example.

class IsRotationInvariant {α β : Type} [Inhabited β] (aug : Augmentation α β) : Prop where combine_rotate : ∀ (x y : α) (a b c : β), aug.combine y (aug.combine x a b) c = aug.combine x a (aug.combine y b c)

A binary tree whose internal nodes cache an augmentation value.

inductive AugmentedTree (α β : Type) where | empty : AugmentedTree α β | node : AugmentedTree α β → α → β → AugmentedTree α β → AugmentedTree α β deriving Repr, DecidableEq
namespace AugmentedTreevariable {α β : Type} [Inhabited β] (aug : Augmentation α β)

Inorder traversal of the stored values.

@[simp] def keys : AugmentedTree α β → List α | empty => [] | node left key _ right => keys left ++ [key] ++ keys right

The cached augmentation at the root.

@[simp] def storedAug : AugmentedTree α β → β | empty => aug.base | node _ _ a _ => a

The mathematically correct augmentation computed from the children.

@[simp] def realAug : AugmentedTree α β → β | empty => aug.base | node left key _ right => aug.combine key (realAug left) (realAug right)

Every cached augmentation agrees with the mathematically correct one.

@[simp] def WellAugmented : AugmentedTree α β → Prop | empty => True | node left key a right => WellAugmented left ∧ WellAugmented right ∧ a = realAug aug (node left key a right)

Recompute every cached augmentation from the children upward.

@[simp] def recompute : AugmentedTree α β → AugmentedTree α β | empty => empty | node left key _ right => let left' := recompute left let right' := recompute right node left' key (aug.combine key (realAug aug left') (realAug aug right')) right'

Left rotation with local augmentation recomputation.

@[simp] def rotateLeft : AugmentedTree α β → AugmentedTree α β | node a x _ (node b y _ c) => let left' := node a x (aug.combine x (realAug aug a) (realAug aug b)) b node left' y (aug.combine y (realAug aug left') (realAug aug c)) c | t => t

Right rotation with local augmentation recomputation.

@[simp] def rotateRight : AugmentedTree α β → AugmentedTree α β | node (node a x _ b) y _ c => let right' := node b y (aug.combine y (realAug aug b) (realAug aug c)) c node a x (aug.combine x (realAug aug a) (realAug aug right')) right' | t => t

Recomputing preserves the inorder key sequence.

theorem keys_recompute (t : AugmentedTree α β) : keys (recompute aug t) = keys t := by induction t with | empty => rfl | node left key _ right ihLeft ihRight => simp [recompute, keys, ihLeft, ihRight]

Recomputing preserves the mathematical augmentation.

theorem realAug_recompute (t : AugmentedTree α β) : realAug aug (recompute aug t) = realAug aug t := by induction t with | empty => rfl | node left key _ right ihLeft ihRight => simp [recompute, realAug, ihLeft, ihRight]

Recomputing establishes the well-augmented invariant.

theorem recompute_wellAugmented (t : AugmentedTree α β) : WellAugmented aug (recompute aug t) := by induction t with | empty => trivial | node left key _ right ihLeft ihRight => simp [recompute, WellAugmented, realAug_recompute, ihLeft, ihRight]

A well-augmented tree has a correct root augmentation.

theorem storedAug_eq_realAug_of_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : storedAug aug t = realAug aug t := by cases t with | empty => rfl | node left key a right => exact h.2.2

Left rotation preserves the inorder key sequence.

theorem keys_rotateLeft (t : AugmentedTree α β) : keys (rotateLeft aug t) = keys t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => simp [rotateLeft, keys, List.append_assoc]

Right rotation preserves the inorder key sequence.

theorem keys_rotateRight (t : AugmentedTree α β) : keys (rotateRight aug t) = keys t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => simp [rotateRight, keys, List.append_assoc]

Left rotation preserves the mathematical augmentation.

theorem realAug_rotateLeft (t : AugmentedTree α β) [IsRotationInvariant aug] : realAug aug (rotateLeft aug t) = realAug aug t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => simp [rotateLeft, realAug] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves the mathematical augmentation.

theorem realAug_rotateRight (t : AugmentedTree α β) [IsRotationInvariant aug] : realAug aug (rotateRight aug t) = realAug aug t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => simp [rotateRight, realAug] rw [IsRotationInvariant.combine_rotate]

Left rotation preserves the cached root augmentation of a well-augmented tree.

theorem storedAug_rotateLeft_of_wellAugmented {t : AugmentedTree α β} [IsRotationInvariant aug] (h : WellAugmented aug t) : storedAug aug (rotateLeft aug t) = storedAug aug t := by cases t with | empty => rfl | node a x _ right => cases right with | empty => rfl | node b y _ c => rcases h with ⟨_ha, hRight, hSize⟩ simp [rotateLeft, storedAug, realAug, hSize] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves the cached root augmentation of a well-augmented tree.

theorem storedAug_rotateRight_of_wellAugmented {t : AugmentedTree α β} [IsRotationInvariant aug] (h : WellAugmented aug t) : storedAug aug (rotateRight aug t) = storedAug aug t := by cases t with | empty => rfl | node left y _ c => cases left with | empty => rfl | node a x _ b => rcases h with ⟨hLeft, _hc, hSize⟩ simp [rotateRight, storedAug, realAug, hSize] rw [IsRotationInvariant.combine_rotate]

Left rotation preserves the well-augmented invariant.

theorem rotateLeft_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (rotateLeft aug t) := by cases t with | empty => exact h | node a x _ right => cases right with | empty => simpa [rotateLeft] using h | node b y _ c => rcases h with ⟨ha, hRight, hSize⟩ rcases hRight with ⟨hb, hc, hRightSize⟩ simp [rotateLeft, WellAugmented, realAug, ha, hb, hc]

Right rotation preserves the well-augmented invariant.

theorem rotateRight_wellAugmented {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (rotateRight aug t) := by cases t with | empty => exact h | node left y _ c => cases left with | empty => simpa [rotateRight] using h | node a x _ b => rcases h with ⟨hLeft, hc, hSize⟩ rcases hLeft with ⟨ha, hb, hLeftSize⟩ simp [rotateRight, WellAugmented, realAug, ha, hb, hc]
Generic BST insertion and the general augmentation theorem (CLRS 14.1)

BST insertion by a Boolean comparison lt, recomputing the cached augmentation locally at each node. Generic over any augmentation.

def insert (lt : α → α → Bool) (x : α) : AugmentedTree α β → AugmentedTree α β | empty => node empty x (aug.combine x aug.base aug.base) empty | node l k _ r => if lt x k then let l' := insert lt x l node l' k (aug.combine k (realAug aug l') (realAug aug r)) r else if lt k x then let r' := insert lt x r node l k (aug.combine k (realAug aug l) (realAug aug r')) r' else node l k (aug.combine k (realAug aug l) (realAug aug r)) r

Insertion adds exactly the inserted key to the inorder key multiset.

theorem mem_keys_insert (lt : α → α → Bool) (x y : α) (t : AugmentedTree α β) : y ∈ keys (insert aug lt x t) → y = x ∨ y ∈ keys t := by induction t with | empty => simp [insert, keys] | node l k a r ihl ihr => simp only [insert] split · simp only [keys, List.mem_append, List.mem_singleton] rintro ((h | h) | h) · rcases ihl h with h' | h' <;> tauto · tauto · tauto · split · simp only [keys, List.mem_append, List.mem_singleton] rintro ((h | h) | h) · tauto · tauto · rcases ihr h with h' | h' <;> tauto · simp only [keys, List.mem_append, List.mem_singleton] tauto

Generic insertion preserves the well-augmented invariant.

theorem insert_wellAugmented (lt : α → α → Bool) (x : α) {t : AugmentedTree α β} (h : WellAugmented aug t) : WellAugmented aug (insert aug lt x t) := by induction t with | empty => simp [insert, WellAugmented, realAug] | node l k a r ihl ihr => rcases h with ⟨hl, hr, _ha⟩ simp only [insert] split · exact ⟨ihl hl, hr, by simp [realAug]⟩ · split · exact ⟨hl, ihr hr, by simp [realAug]⟩ · exact ⟨hl, hr, by simp [realAug]⟩

CLRS Theorem 14.1 (maintainability of augmentations). For any locally-computable, rotation-invariant augmentation, every structural primitive used by red-black insertion and deletion — left and right rotation, subtree recomputation, and BST insertion — preserves the inorder key sequence and the mathematical augmentation, and preserves (or re-establishes) the WellAugmented invariant. Hence the augmentation can be maintained through the red-black operations.

theorem augmentation_theorem [IsRotationInvariant aug] : (∀ t : AugmentedTree α β, WellAugmented aug t → WellAugmented aug (rotateLeft aug t)) ∧ (∀ t : AugmentedTree α β, WellAugmented aug t → WellAugmented aug (rotateRight aug t)) ∧ (∀ t : AugmentedTree α β, keys (rotateLeft aug t) = keys t) ∧ (∀ t : AugmentedTree α β, keys (rotateRight aug t) = keys t) ∧ (∀ t : AugmentedTree α β, realAug aug (rotateLeft aug t) = realAug aug t) ∧ (∀ t : AugmentedTree α β, realAug aug (rotateRight aug t) = realAug aug t) ∧ (∀ t : AugmentedTree α β, WellAugmented aug (recompute aug t)) ∧ (∀ (lt : α → α → Bool) (x : α) (t : AugmentedTree α β), WellAugmented aug t → WellAugmented aug (insert aug lt x t)) := ⟨fun _ h => rotateLeft_wellAugmented aug h, fun _ h => rotateRight_wellAugmented aug h, fun t => keys_rotateLeft aug t, fun t => keys_rotateRight aug t, fun t => realAug_rotateLeft aug t, fun t => realAug_rotateRight aug t, fun t => recompute_wellAugmented aug t, fun lt x _ h => insert_wellAugmented aug lt x h⟩
end AugmentedTree
Size augmentation (order-statistic trees as an instance)

The order-statistic augmentation of Section 14.1 — caching each subtree's node count — is an instance of the same generic framework, demonstrating CLRS Theorem 14.1 for a second concrete field alongside interval trees' max-high.

The subtree-size augmentation: the cached value is the number of nodes.

def sizeAug (α : Type) : Augmentation α Nat := ⟨0, fun _ l r => 1 + l + r⟩
instance (α : Type) : IsRotationInvariant (sizeAug α) where combine_rotate x y a b c := by simp [sizeAug]; omega

The size augmentation's mathematical value is exactly the node count.

theorem realAug_sizeAug_eq_length {α : Type} (t : AugmentedTree α Nat) : AugmentedTree.realAug (sizeAug α) t = (AugmentedTree.keys t).length := by induction t with | empty => rfl | node l k a r ihl ihr => simp only [AugmentedTree.realAug] rw [ihl, ihr] simp only [sizeAug, AugmentedTree.keys, List.length_append, List.length_cons, List.length_nil] omega

Interval trees

A closed interval of natural numbers.

def Interval := Nat × Nat
def Interval.low (i : Interval) : Nat := i.1def Interval.high (i : Interval) : Nat := i.2def Interval.overlaps (i j : Interval) : Bool := i.low ≤ j.high && j.low ≤ i.high@[simp] theorem Interval.overlaps_iff {i j : Interval} : Interval.overlaps i j = true ↔ i.low ≤ j.high ∧ j.low ≤ i.high := by simp [Interval.overlaps]

Interval trees are augmented trees whose node value is an interval and whose augmentation is the maximum high endpoint in the subtree.

abbrev IntervalTree := AugmentedTree Interval Nat
def IntervalTree.maxHighAug : Augmentation Interval Nat := ⟨0, fun i l r => max i.high (max l r)⟩instance : IsRotationInvariant IntervalTree.maxHighAug where combine_rotate x y a b c := by simp [IntervalTree.maxHighAug] ac_rflnamespace IntervalTree

Inorder list of intervals.

def keys : IntervalTree → List Interval := AugmentedTree.keys

Cached maximum high endpoint at the root.

def storedMaxHigh : IntervalTree → Nat := AugmentedTree.storedAug maxHighAug

Mathematical maximum high endpoint in the subtree.

def realMaxHigh : IntervalTree → Nat := AugmentedTree.realAug maxHighAug

The max-high augmentation invariant.

def WellAugmented : IntervalTree → Prop := AugmentedTree.WellAugmented maxHighAug

Recompute every cached max-high field.

Left rotation with local max-high recomputation.

Right rotation with local max-high recomputation.

Every interval in the tree has low endpoint at most x.

def allLowLE (t : IntervalTree) (x : Nat) : Prop := ∀ i ∈ keys t, i.low ≤ x

Every interval in the tree has low endpoint at least x.

def allLowGE (t : IntervalTree) (x : Nat) : Prop := ∀ i ∈ keys t, x ≤ i.low

Binary-search-tree ordering by interval low endpoint.

def IsBST : IntervalTree → Prop | AugmentedTree.empty => True | AugmentedTree.node left int _ right => IsBST left ∧ IsBST right ∧ allLowLE left int.low ∧ allLowGE right int.low

Boolean emptiness test for interval trees.

def isEmpty : IntervalTree → Bool | AugmentedTree.empty => true | AugmentedTree.node _ _ _ _ => false

Decision to recurse into the left subtree during interval search.

def goLeft (left : IntervalTree) (q : Interval) : Bool := !isEmpty left && decide (storedMaxHigh left ≥ q.low)

The executable interval-search algorithm from CLRS.

def intervalSearch? : IntervalTree → Interval → Option Interval | AugmentedTree.empty, _ => none | AugmentedTree.node left int _ right, q => if Interval.overlaps int q then some int else if goLeft left q then intervalSearch? left q else intervalSearch? right q

Does the tree contain an interval overlapping the query?

def hasOverlap (t : IntervalTree) (q : Interval) : Prop := ∃ i ∈ keys t, Interval.overlaps i q
end IntervalTreenamespace IntervalTree@[simp] theorem keys_empty : keys AugmentedTree.empty = [] := by simp [keys]@[simp] theorem keys_node {left right : IntervalTree} {int : Interval} {mx : Nat} : keys (AugmentedTree.node left int mx right) = keys left ++ [int] ++ keys right := by simp [keys]@[simp] theorem isEmpty_empty : isEmpty AugmentedTree.empty = true := by rfl@[simp] theorem isEmpty_node {left right : IntervalTree} {int : Interval} {mx : Nat} : isEmpty (AugmentedTree.node left int mx right) = false := by rfl

A true goLeft condition means the left subtree is non-empty and its max-high is at least the query low.

theorem goLeft_true {left : IntervalTree} {q : Interval} (h : goLeft left q = true) : left ≠ AugmentedTree.empty ∧ storedMaxHigh left ≥ q.low := by simp [goLeft, Bool.and_eq_true] at h rcases h with ⟨hne, hmax⟩ constructor · cases left with | empty => simp at hne | node => simp · exact hmax

A false goLeft condition means the left subtree is empty or its max-high is below the query low.

theorem goLeft_false {left : IntervalTree} {q : Interval} (h : goLeft left q = false) : left = AugmentedTree.empty ∨ storedMaxHigh left < q.low := by simp [goLeft] at h cases left with | empty => left; rfl | node left int mx right => right have hmax : decide (storedMaxHigh (AugmentedTree.node left int mx right) ≥ q.low) = false := by simpa using h simpa using hmax

Recomputing cached max-high fields preserves the inorder key sequence.

theorem keys_recompute (t : IntervalTree) : keys (recompute t) = keys t := AugmentedTree.keys_recompute maxHighAug t

Recomputing cached max-high fields preserves the mathematical max-high.

theorem realMaxHigh_recompute (t : IntervalTree) : realMaxHigh (recompute t) = realMaxHigh t := AugmentedTree.realAug_recompute maxHighAug t

Recomputing cached max-high fields establishes the augmentation invariant.

theorem recompute_wellAugmented (t : IntervalTree) : WellAugmented (recompute t) := AugmentedTree.recompute_wellAugmented maxHighAug t

A well-augmented interval tree has a correct root max-high field.

theorem storedMaxHigh_eq_realMaxHigh_of_wellAugmented {t : IntervalTree} (h : WellAugmented t) : storedMaxHigh t = realMaxHigh t := AugmentedTree.storedAug_eq_realAug_of_wellAugmented maxHighAug h

Left rotation preserves the well-augmented invariant.

theorem rotateLeft_wellAugmented {t : IntervalTree} (h : WellAugmented t) : WellAugmented (rotateLeft t) := AugmentedTree.rotateLeft_wellAugmented maxHighAug h

Right rotation preserves the well-augmented invariant.

Membership in keys respects the inorder list membership relation.

@[simp] theorem mem_keys {t : IntervalTree} {i : Interval} : i ∈ keys t ↔ i ∈ AugmentedTree.keys t := by rfl

Every stored high endpoint is bounded by the real max-high.

theorem high_le_realMaxHigh {t : IntervalTree} {i : Interval} (hi : i ∈ keys t) : i.high ≤ realMaxHigh t := by induction t with | empty => simp [keys] at hi | node left int _ right ihLeft ihRight => simp [keys] at hi rcases hi with hi | hi | hi · exact le_trans (ihLeft hi) (by simp [realMaxHigh, AugmentedTree.realAug, maxHighAug]) · simp [hi, realMaxHigh, AugmentedTree.realAug, maxHighAug] · exact le_trans (ihRight hi) (by simp [realMaxHigh, AugmentedTree.realAug, maxHighAug])

If the mathematical max-high of a subtree is below the query low, the subtree contains no overlap.

theorem noOverlap_of_realMaxHigh_lt {t : IntervalTree} {q : Interval} (h : realMaxHigh t < q.low) : ¬ ∃ i ∈ keys t, Interval.overlaps i q := by rintro ⟨i, hi, hov⟩ rw [Interval.overlaps_iff] at hov have hiHigh := high_le_realMaxHigh hi have : q.low ≤ i.high := hov.2 linarith

A tree whose realMaxHigh is at least a positive bound contains a member whose high is at least that bound.

private theorem exists_mem_high_ge {t : IntervalTree} {x : Nat} (hx : x > 0) (h : realMaxHigh t ≥ x) : ∃ i ∈ keys t, i.high ≥ x := by induction t with | empty => simp [realMaxHigh, AugmentedTree.realAug, maxHighAug] at h omega | node left int _ right ihLeft ihRight => simp [realMaxHigh, AugmentedTree.realAug, maxHighAug] at h by_cases hi : int.high ≥ x · use int; simp [hi, keys_node] · have : realMaxHigh left ≥ x ∨ realMaxHigh right ≥ x := by simp [realMaxHigh, maxHighAug] at h ⊢ omega rcases this with h' | h' · obtain ⟨k, hk, hkHigh⟩ := ihLeft h' use k; simp [hk, hkHigh, keys_node] · obtain ⟨k, hk, hkHigh⟩ := ihRight h' use k; simp [hk, hkHigh, keys_node]

If the left subtree is non-empty, its max-high is at least the query low, and the current interval does not overlap, then any overlap in the right subtree forces an overlap in the left subtree. This is the key pruning invariant for interval search.

theorem overlap_left_of_right_overlap {left right : IntervalTree} {int q : Interval} (mx : Nat) (hB : IsBST (AugmentedTree.node left int mx right)) (hmax : realMaxHigh left ≥ q.low) (hcur : ¬ Interval.overlaps int q) (hright : ∃ j ∈ keys right, Interval.overlaps j q) : ∃ i ∈ keys left, Interval.overlaps i q := by rcases hright with ⟨j, hj, hov⟩ rw [Interval.overlaps_iff] at hov rcases hB with ⟨_hBL, _hBR, hLeftLE, hRightGE⟩ have h1 : int.low ≤ j.low := hRightGE j hj have h2 : j.low ≤ q.high := hov.1 have h3 : int.low ≤ q.high := by linarith have h4 : int.high < q.low := by by_contra h' have hov : Interval.overlaps int q = true := by rw [Interval.overlaps_iff] exact ⟨h3, by omega⟩ exact hcur hov have h5 : ∃ i ∈ keys left, i.high ≥ q.low := by have hx : q.low > 0 := by omega exact exists_mem_high_ge hx hmax rcases h5 with ⟨i, hi, hiHigh⟩ have h6 : i.low ≤ int.low := hLeftLE i hi have h7 : i.low ≤ q.high := by linarith use i, hi rw [Interval.overlaps_iff] exact ⟨h7, hiHigh⟩
end IntervalTreenamespace IntervalTree

hasOverlap distributes over a node in the obvious way.

@[simp] theorem hasOverlap_node {left right : IntervalTree} {int : Interval} {mx : Nat} {q : Interval} : hasOverlap (AugmentedTree.node left int mx right) q ↔ Interval.overlaps int q = true ∨ hasOverlap left q ∨ hasOverlap right q := by simp [hasOverlap, keys_node, Interval.overlaps_iff] constructor · rintro ⟨i, (hi | rfl | hi), hlow, hhigh⟩ · right; left; use i · left; exact ⟨hlow, hhigh⟩ · right; right; use i · rintro (⟨hlow, hhigh⟩ | ⟨i, hi, hlow, hhigh⟩ | ⟨i, hi, hlow, hhigh⟩) · use int; simp; exact ⟨hlow, hhigh⟩ · use i; simp [hi]; exact ⟨hlow, hhigh⟩ · use i; simp [hi]; exact ⟨hlow, hhigh⟩

The executable interval search returns only intervals that are in the tree and overlap the query.

theorem intervalSearch?_some_overlap {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) (i : Interval) : intervalSearch? t q = some i → i ∈ keys t ∧ Interval.overlaps i q := by induction t with | empty => simp [intervalSearch?] | node left int _ right ihLeft ihRight => intro h rcases hB with ⟨hBL, hBR, _hLeftLE, _hRightGE⟩ rcases hW with ⟨hWL, hWR, _hMax⟩ unfold intervalSearch? at h by_cases hO : Interval.overlaps int q = true · -- current interval overlaps simp [hO] at h cases h with | refl => constructor · simp [keys_node] · simp [hO] · -- current interval does not overlap by_cases hL : goLeft left q = true · -- go left simp [hO, hL] at h have hIh := ihLeft hBL hWL h constructor · simp [keys_node, hIh.1] · exact hIh.2 · -- go right simp [hO, hL] at h have hIh := ihRight hBR hWR h constructor · simp [keys_node, hIh.1] · exact hIh.2

If the executable interval search returns none, no interval in the tree overlaps the query.

theorem intervalSearch?_none_noOverlap {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) : intervalSearch? t q = none → ¬ hasOverlap t q := by induction t with | empty => simp [intervalSearch?, hasOverlap] | node left int mx right ihLeft ihRight => intro h rcases hB with ⟨hBL, hBR, hLeftLE, hRightGE⟩ rcases hW with ⟨hWL, hWR, _hMax⟩ unfold intervalSearch? at h by_cases hO : Interval.overlaps int q = true · -- current overlaps, but returned none: impossible simp [hO] at h · -- current interval does not overlap by_cases hL : goLeft left q = true · -- went left and got none simp [hO, hL] at h have hNoLeft : ¬ hasOverlap left q := ihLeft hBL hWL h intro hOv rcases hOv with ⟨j, hj, hov⟩ simp [keys_node] at hj rcases hj with (hjLeft | hjEq | hjRight) · -- overlap in left subtree exact hNoLeft ⟨j, hjLeft, hov⟩ · -- overlap with current interval rw [hjEq] at hov exact hO hov · -- overlap in right subtree forces one in left subtree have hRightEx : ∃ j ∈ keys right, Interval.overlaps j q := ⟨j, hjRight, hov⟩ have hL' := goLeft_true hL have hLeftOverlap := overlap_left_of_right_overlap mx ⟨hBL, hBR, hLeftLE, hRightGE⟩ (by rw [← storedMaxHigh_eq_realMaxHigh_of_wellAugmented hWL]; exact hL'.2) hO hRightEx rcases hLeftOverlap with ⟨k, hk, hkov⟩ exact hNoLeft ⟨k, hk, hkov⟩ · -- went right and got none have hL_false : goLeft left q = false := by simp [hL] simp [hO, hL] at h have hNoRight : ¬ hasOverlap right q := ihRight hBR hWR h intro hOv rcases hOv with ⟨j, hj, hov⟩ simp [keys_node] at hj rcases hj with (hjLeft | hjEq | hjRight) · -- no overlap possible in left subtree have hNoLeft : ¬ hasOverlap left q := by rcases goLeft_false hL_false with hEmpty | hmax · simp [hEmpty, hasOverlap] · have hmax' : realMaxHigh left < q.low := by rw [← storedMaxHigh_eq_realMaxHigh_of_wellAugmented hWL] exact hmax intro hOv' exact noOverlap_of_realMaxHigh_lt hmax' hOv' exact hNoLeft ⟨j, hjLeft, hov⟩ · -- overlap with current interval rw [hjEq] at hov exact hO hov · -- overlap in right subtree exact hNoRight ⟨j, hjRight, hov⟩

Combined correctness specification for interval search.

theorem intervalSearch?_spec {t : IntervalTree} (hB : IsBST t) (hW : WellAugmented t) (q : Interval) : (intervalSearch? t q = none ↔ ¬ hasOverlap t q) ∧ (∀ i, intervalSearch? t q = some i → i ∈ keys t ∧ Interval.overlaps i q) := by constructor · constructor · exact intervalSearch?_none_noOverlap hB hW q · intro hNoOverlap by_contra h have : intervalSearch? t q ≠ none := by simp [h] rcases Option.ne_none_iff_exists'.mp this with ⟨i, hi⟩ have := intervalSearch?_some_overlap hB hW q i hi exact hNoOverlap ⟨i, this.1, this.2⟩ · exact intervalSearch?_some_overlap hB hW q
end IntervalTree
Red-black bridge: maintaining an augmentation through Chapter 13 rotations

CLRS Theorem 14.1 in the red-black setting: a locally-computable augmentation can be maintained through the structural primitives used by red-black insertion and deletion. Chapter 13's red-black rotations and root recoloring are exactly those primitives. We show that any rotation-invariant augmentation's value is preserved by these operations — so a locally-recomputed cached field stays correct — reusing the same IsRotationInvariant law as the generic framework.

Note: red-black rotations are shape-restoring, not shape-preserving: they are applied mid-fixup where RedBlackShape is temporarily broken, so we do not (and cannot) claim a single rotation preserves RedBlackShape. RedBlackShape maintenance across a full insertion is Chapter 13's RBTree.redBlackShape_insert; the new content here is that the augmentation rides along invariantly under the same rotations and recoloring.

namespace RBBridgeopen CLRS.Chapter13

Inorder key list of a Chapter 13 red-black tree.

def rbKeys : RBTree → List Nat | .empty => [] | .node _ l k r => rbKeys l ++ [k] ++ rbKeys r

Semantic value of an augmentation on a red-black tree, computed from each key and its children (independent of node colors).

def rbRealAug {β : Type} [Inhabited β] (aug : Augmentation Nat β) : RBTree → β | .empty => aug.base | .node _ l k r => aug.combine k (rbRealAug aug l) (rbRealAug aug r)

Left rotation preserves the inorder key list.

theorem rbKeys_rotateLeft (t : RBTree) : rbKeys (RBTree.rotateLeft t) = rbKeys t := by cases t with | empty => rfl | node color a x right => cases right with | empty => rfl | node rc b y c => simp [RBTree.rotateLeft, rbKeys, List.append_assoc]

Right rotation preserves the inorder key list.

theorem rbKeys_rotateRight (t : RBTree) : rbKeys (RBTree.rotateRight t) = rbKeys t := by cases t with | empty => rfl | node color left y c => cases left with | empty => rfl | node lc a x b => simp [RBTree.rotateRight, rbKeys, List.append_assoc]

Left rotation preserves any rotation-invariant augmentation's value.

theorem rbRealAug_rotateLeft {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] (t : RBTree) : rbRealAug aug (RBTree.rotateLeft t) = rbRealAug aug t := by cases t with | empty => rfl | node color a x right => cases right with | empty => rfl | node rc b y c => simp only [RBTree.rotateLeft, rbRealAug] rw [IsRotationInvariant.combine_rotate]

Right rotation preserves any rotation-invariant augmentation's value.

theorem rbRealAug_rotateRight {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] (t : RBTree) : rbRealAug aug (RBTree.rotateRight t) = rbRealAug aug t := by cases t with | empty => rfl | node color left y c => cases left with | empty => rfl | node lc a x b => simp only [RBTree.rotateRight, rbRealAug] rw [IsRotationInvariant.combine_rotate]

Root recoloring preserves the augmentation value.

theorem rbRealAug_repaintRoot {β : Type} [Inhabited β] (aug : Augmentation Nat β) (c : Color) (t : RBTree) : rbRealAug aug (RBTree.repaintRoot c t) = rbRealAug aug t := by cases t <;> simp [RBTree.repaintRoot, rbRealAug]

Red-black bridge (CLRS Theorem 14.1, red-black primitives). Every structural primitive used by red-black insertion and deletion — left and right rotation and root recoloring — preserves both the inorder key sequence and any rotation-invariant augmentation's value. Hence the augmentation can be maintained through the red-black operations by local recomputation, exactly as in the generic framework.

theorem rb_augmentation_bridge {β : Type} [Inhabited β] (aug : Augmentation Nat β) [IsRotationInvariant aug] : (∀ t, rbKeys (RBTree.rotateLeft t) = rbKeys t) ∧ (∀ t, rbKeys (RBTree.rotateRight t) = rbKeys t) ∧ (∀ t, rbRealAug aug (RBTree.rotateLeft t) = rbRealAug aug t) ∧ (∀ t, rbRealAug aug (RBTree.rotateRight t) = rbRealAug aug t) ∧ (∀ (c : Color) (t), rbRealAug aug (RBTree.repaintRoot c t) = rbRealAug aug t) := ⟨rbKeys_rotateLeft, rbKeys_rotateRight, rbRealAug_rotateLeft aug, rbRealAug_rotateRight aug, fun c t => rbRealAug_repaintRoot aug c t⟩

The size augmentation's value on a red-black tree is its node count.

theorem rbRealAug_sizeAug_eq_length (t : RBTree) : rbRealAug (sizeAug Nat) t = (rbKeys t).length := by induction t with | empty => rfl | node c l k r ihl ihr => simp only [rbRealAug, rbKeys] rw [ihl, ihr] simp only [sizeAug, List.length_append, List.length_cons, List.length_nil] omega
end RBBridge
General augmentation interface: an arbitrary augmentation through

executable red-black insertion

This section closes the "stored-field refinement" gap noted above, at the generic level. Section 14.1's OSRBTree threaded only the concrete subtree-size augmentation through Chapter 13's executable red-black insertion, in a bespoke, size-specific type. Here we thread an arbitrary Augmentation through the same Okasaki-style balancer, so both the order-statistic (size) and interval (max-high) augmentations are recovered as instances of a single generic interface.

The augmented red-black tree AugmentedRBTree caches, at every internal node, a node colour (reusing Chapter 13's CLRS.­Chapter13.­Color) and an augmentation value of type β. Every reconstructed node is built by the smart constructor AugmentedRBTree.mk, which recomputes the cached augmentation from its children via aug.combine. Two bridges connect this to the existing development:

  • AugmentedRBTree.wellAugmented_insert: the augmentation invariant survives balancing — inserting into a well-augmented tree yields a well-augmented tree, for any augmentation (CLRS 14.1 maintained through RB-INSERT).

  • AugmentedRBTree.toRB_insert: erasing the augmentation field commutes with insertion, so (for Nat keys) the augmented insert refines the executable Chapter 13 CLRS.­Chapter13.­RBTree.­insert exactly, transferring its shape and membership theorems.

The size and max-high fields are then recovered as instances via AugmentedRBTree.sizeAug_wellAugmented_insert and AugmentedRBTree.maxHighAug_wellAugmented_insert.

open CLRS.Chapter13 (Color RBTree)

A red-black tree augmented with a cached value of type β at every internal node. Each node stores a colour (reusing Chapter 13's CLRS.­Chapter13.­Color), a key of type α, a cached augmentation of type β, and two subtrees. This is the colour-carrying refinement of AugmentedTree, adding the field needed to run the Chapter 13 insertion balancer generically.

inductive AugmentedRBTree (α β : Type) where | empty : AugmentedRBTree α β | node : Color → AugmentedRBTree α β → α → β → AugmentedRBTree α β → AugmentedRBTree α β deriving Repr, DecidableEq
namespace AugmentedRBTreesection Genericvariable {α β : Type} [Inhabited β] (aug : Augmentation α β)

Inorder traversal of the keys, ignoring colours and cached augmentations.

def keys : AugmentedRBTree α β → List α | empty => [] | node _ l k _ r => keys l ++ [k] ++ keys r

The cached augmentation stored at the root; the empty tree uses aug.base.

def storedAug : AugmentedRBTree α β → β | empty => aug.base | node _ _ _ a _ => a

The mathematically correct augmentation, recomputed from the children.

def realAug : AugmentedRBTree α β → β | empty => aug.base | node _ l k _ r => aug.combine k (realAug l) (realAug r)

Every cached augmentation agrees with the recomputed one.

def WellAugmented : AugmentedRBTree α β → Prop | empty => True | node _ l k a r => WellAugmented l ∧ WellAugmented r ∧ a = aug.combine k (realAug aug l) (realAug aug r)

Smart constructor that recomputes the cached augmentation from the children. Every node produced by the red-black operations below is built with mk, which is why the WellAugmented invariant is preserved automatically.

def mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := node c l k (aug.combine k (realAug aug l) (realAug aug r)) r

mk recomputes the augmentation correctly.

theorem realAug_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : realAug aug (mk aug c l k r) = aug.combine k (realAug aug l) (realAug aug r) := rfl

mk preserves the inorder key sequence.

theorem keys_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : keys (mk aug c l k r) = keys l ++ [k] ++ keys r := rfl

The cached root augmentation of a mk node is the recomputed value.

theorem storedAug_mk (c : Color) (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : storedAug aug (mk aug c l k r) = aug.combine k (realAug aug l) (realAug aug r) := rfl

A mk node is well-augmented whenever both children are.

theorem wellAugmented_mk {c : Color} {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (mk aug c l k r) := ⟨hl, hr, rfl⟩

A well-augmented tree has a correct root augmentation.

theorem storedAug_eq_realAug_of_wellAugmented {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug t = realAug aug t := by cases t with | empty => rfl | node c l k a r => exact h.2.2
Executable red-black operations with augmentation recomputation

Repaint the root black, keeping the cached augmentation fields.

def repaintBlack : AugmentedRBTree α β → AugmentedRBTree α β | empty => empty | node _ l k a r => node Color.black l k a r

Okasaki-style rebalance after insertion on the left child, recomputing augmentations. Mirrors CLRS.­Chapter13.­RBTree.­balanceLeft.

def balanceLeft (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match l with | node Color.red (node Color.red a w _ b) x _ c => mk aug Color.red (mk aug Color.black a w b) x (mk aug Color.black c y r) | node Color.red a w _ (node Color.red b x _ c) => mk aug Color.red (mk aug Color.black a w b) x (mk aug Color.black c y r) | _ => mk aug Color.black l y r

Okasaki-style rebalance after insertion on the right child, recomputing augmentations. Mirrors CLRS.­Chapter13.­RBTree.­balanceRight.

def balanceRight (l : AugmentedRBTree α β) (y : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | node Color.red (node Color.red b x _ c) y' _ d => mk aug Color.red (mk aug Color.black l y b) x (mk aug Color.black c y' d) | node Color.red b x _ (node Color.red c y' _ d) => mk aug Color.red (mk aug Color.black l y b) x (mk aug Color.black c y' d) | _ => mk aug Color.black l y r

Insertion fixup: recurse down by the Boolean comparison lt, rebuilding and rebalancing with augmentation recomputation on the way back up. Mirrors CLRS.­Chapter13.­RBTree.­insertFixup.

def insertFixup (lt : α → α → Bool) (x : α) : AugmentedRBTree α β → AugmentedRBTree α β | empty => mk aug Color.red empty x empty | node c l y a r => if lt x y then if c = Color.black then balanceLeft aug (insertFixup lt x l) y r else mk aug Color.red (insertFixup lt x l) y r else if lt y x then if c = Color.black then balanceRight aug l y (insertFixup lt x r) else mk aug Color.red l y (insertFixup lt x r) else node c l y a r

Insert a key and repaint the root black.

def insert (lt : α → α → Bool) (x : α) (t : AugmentedRBTree α β) : AugmentedRBTree α β := repaintBlack (insertFixup aug lt x t)
The augmentation invariant survives balancing (CLRS 14.1 through RB-INSERT)

Repainting the root black preserves the WellAugmented invariant.

theorem wellAugmented_repaintBlack {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (repaintBlack t) := by cases t with | empty => trivial | node c l k a r => exact ⟨h.1, h.2.1, h.2.2⟩

balanceLeft preserves the WellAugmented invariant.

theorem wellAugmented_balanceLeft {l : AugmentedRBTree α β} {y : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (balanceLeft aug l y r) := by unfold balanceLeft split · obtain ⟨⟨ha, hb, _⟩, hc, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) (wellAugmented_mk aug hc hr) · obtain ⟨ha, ⟨hb, hc, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) (wellAugmented_mk aug hc hr) · exact wellAugmented_mk aug hl hr

balanceRight preserves the WellAugmented invariant.

theorem wellAugmented_balanceRight {l : AugmentedRBTree α β} {y : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (balanceRight aug l y r) := by unfold balanceRight split · obtain ⟨⟨hb, hc, _⟩, hd, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hb) (wellAugmented_mk aug hc hd) · obtain ⟨hb, ⟨hc, hd, _⟩, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hb) (wellAugmented_mk aug hc hd) · exact wellAugmented_mk aug hl hr

insertFixup preserves the WellAugmented invariant.

theorem wellAugmented_insertFixup (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (insertFixup aug lt x t) := by induction t with | empty => simp only [insertFixup] exact wellAugmented_mk aug (by trivial) (by trivial) | node c l y a r ihl ihr => have hl : WellAugmented aug l := h.1 have hr : WellAugmented aug r := h.2.1 simp only [insertFixup] split · split · exact wellAugmented_balanceLeft aug (ihl hl) hr · exact wellAugmented_mk aug (ihl hl) hr · split · split · exact wellAugmented_balanceRight aug hl (ihr hr) · exact wellAugmented_mk aug hl (ihr hr) · exact h

Augmentation invariant through executable insertion (CLRS 14.1 through RB-INSERT). Inserting a key into a well-augmented augmented red-black tree produces a well-augmented tree: every cached augmentation field remains correct after the red-black rebalancing, for any Augmentation. This generalizes OSRBTree.wellSized_insert from the size field to an arbitrary augmentation.

theorem wellAugmented_insert (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (insert aug lt x t) := by unfold insert exact wellAugmented_repaintBlack aug (wellAugmented_insertFixup aug lt x h)

After insertion the cached root augmentation equals the recomputed value.

theorem storedAug_insert (lt : α → α → Bool) (x : α) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug (insert aug lt x t) = realAug aug (insert aug lt x t) := storedAug_eq_realAug_of_wellAugmented aug (wellAugmented_insert aug lt x h)
Executable red-black deletion with augmentation recomputation

The deletion pipeline mirrors OSRBTree but is generic in α, β, and aug. Every function recomputes augmentations via mk aug.

Repaint with arbitrary color, keeping cached augmentation fields.

def repaintRoot (c : Color) (t : AugmentedRBTree α β) : AugmentedRBTree α β := match t with | empty => empty | node _ l k a r => node c l k a r

Boolean black-root test.

def rootBlack : AugmentedRBTree α β → Bool | empty => true | node c _ _ _ _ => c = Color.black

Deletion re-balancer for a black-deficient left child.

def baldL (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match l with | node Color.red a x _ b => mk aug Color.red (mk aug Color.black a x b) k r | _ => match r with | node Color.black c y _ d => balanceRight aug l k (mk aug Color.red c y d) | node Color.red (node Color.black c y _ d) z _ e => mk aug Color.red (mk aug Color.black l k c) y (balanceRight aug d z (repaintRoot Color.red e)) | _ => mk aug Color.red l k r

Deletion re-balancer for a black-deficient right child.

def baldR (l : AugmentedRBTree α β) (k : α) (r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | node Color.red c y _ d => mk aug Color.red l k (mk aug Color.black c y d) | _ => match l with | node Color.black a x _ b => balanceLeft aug (mk aug Color.red a x b) k r | node Color.red a x _ (node Color.black c y _ d) => mk aug Color.red (balanceLeft aug (repaintRoot Color.red a) x c) y (mk aug Color.black d k r) | _ => mk aug Color.red l k r

Find and remove the minimum key, recomputing augmentations on the way up.

def splitMin [Inhabited α] : AugmentedRBTree α β → α × AugmentedRBTree α β | empty => (default, empty) | node _ empty k _ r => (k, r) | node _ l k _ r => let (m, l') := splitMin l if rootBlack l then (m, baldL aug l' k r) else (m, mk aug Color.red l' k r)

Merge two trees, used when deleting a node with two children.

def join [Inhabited α] (l r : AugmentedRBTree α β) : AugmentedRBTree α β := match r with | empty => l | _ => match l with | empty => r | _ => let (m, r') := splitMin aug r if rootBlack r then baldR aug l m r' else mk aug Color.red l m r'

Recursive deletion.

def del [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) : AugmentedRBTree α β → AugmentedRBTree α β | empty => empty | node _c l y _a r => if lt x y then if rootBlack l then baldL aug (del x lt l) y r else mk aug Color.red (del x lt l) y r else if lt y x then if rootBlack r then baldR aug l y (del x lt r) else mk aug Color.red l y (del x lt r) else join aug l r

Delete a key and repaint the root black.

def delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) (t : AugmentedRBTree α β) : AugmentedRBTree α β := repaintBlack (del aug x lt t)
The augmentation invariant survives deletion
theorem wellAugmented_repaintRoot (c : Color) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (repaintRoot c t) := by cases t with | empty => trivial | node c' l k a r => exact ⟨h.1, h.2.1, h.2.2⟩theorem wellAugmented_baldL {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (baldL aug l k r) := by unfold baldL -- Use case analysis that matches the definition's pattern-matching structure cases l with | empty => cases r with | empty => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_balanceRight aug hl (wellAugmented_mk aug hc hd) | red => cases rl with | empty => exact wellAugmented_mk aug hl hr | node rlc rll rlk rla rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hc) (wellAugmented_balanceRight aug hd (wellAugmented_repaintRoot aug Color.red he)) | red => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | red => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_mk aug ha hb) hr | black => cases r with | empty => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | black => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_balanceRight aug hl (wellAugmented_mk aug hc hd) | red => cases rl with | empty => exact wellAugmented_mk aug hl hr | node rlc rll rlk rla rlr => cases rlc with | black => obtain ⟨⟨hc, hd, _⟩, he, _⟩ := hr exact wellAugmented_mk aug (wellAugmented_mk aug hl hc) (wellAugmented_balanceRight aug hd (wellAugmented_repaintRoot aug Color.red he)) | red => exact wellAugmented_mk aug hl hrtheorem wellAugmented_baldR {l : AugmentedRBTree α β} {k : α} {r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (baldR aug l k r) := by unfold baldR cases r with | empty => cases l with | empty => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_balanceLeft aug (wellAugmented_mk aug ha hb) hr | red => cases lr with | empty => exact wellAugmented_mk aug hl hr | node lrc lrl lrk lra lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_balanceLeft aug (wellAugmented_repaintRoot aug Color.red ha) hc) (wellAugmented_mk aug hd hr) | red => exact wellAugmented_mk aug hl hr | node rc rl rk ra rr => cases rc with | red => obtain ⟨hc, hd, _⟩ := hr exact wellAugmented_mk aug hl (wellAugmented_mk aug hc hd) | black => cases l with | empty => exact wellAugmented_mk aug hl hr | node lc ll lk la lr => cases lc with | black => obtain ⟨ha, hb, _⟩ := hl exact wellAugmented_balanceLeft aug (wellAugmented_mk aug ha hb) hr | red => cases lr with | empty => exact wellAugmented_mk aug hl hr | node lrc lrl lrk lra lrr => cases lrc with | black => obtain ⟨ha, ⟨hc, hd, _⟩, _⟩ := hl exact wellAugmented_mk aug (wellAugmented_balanceLeft aug (wellAugmented_repaintRoot aug Color.red ha) hc) (wellAugmented_mk aug hd hr) | red => exact wellAugmented_mk aug hl hr theorem wellAugmented_splitMin [Inhabited α] {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (splitMin aug t).2 := by induction t with | empty => trivial | node c l k a r ihl => cases l with | empty => exact h.2.1 | node lc ll lk la lr => have hws : WellAugmented aug (splitMin aug (node lc ll lk la lr)).2 := ihl h.1 by_cases hrb : rootBlack (node lc ll lk la lr) = true · have hsp : (splitMin aug (node c (node lc ll lk la lr) k a r)).2 = baldL aug (splitMin aug (node lc ll lk la lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellAugmented_baldL aug hws h.2.1 · have hsp : (splitMin aug (node c (node lc ll lk la lr) k a r)).2 = mk aug Color.red (splitMin aug (node lc ll lk la lr)).2 k r := by simp [splitMin, hrb] rw [hsp] exact wellAugmented_mk aug hws h.2.1theorem wellAugmented_join [Inhabited α] {l r : AugmentedRBTree α β} (hl : WellAugmented aug l) (hr : WellAugmented aug r) : WellAugmented aug (join aug l r) := by unfold join split · exact hl · rename_i hneR split · exact hr · rename_i hneL dsimp by_cases hrb : rootBlack r = true · simp [hrb] exact wellAugmented_baldR aug hl (wellAugmented_splitMin aug hr) · simp [hrb] exact wellAugmented_mk aug hl (wellAugmented_splitMin aug hr)theorem wellAugmented_del [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (del aug x lt t) := by induction t with | empty => exact h | node c l y a r ihl ihr => have hl : WellAugmented aug l := h.1 have hr : WellAugmented aug r := h.2.1 simp only [del] split · split · exact wellAugmented_baldL aug (ihl hl) hr · exact wellAugmented_mk aug (ihl hl) hr · split · split · exact wellAugmented_baldR aug hl (ihr hr) · exact wellAugmented_mk aug hl (ihr hr) · exact wellAugmented_join aug hl hrtheorem wellAugmented_delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : WellAugmented aug (delete aug x lt t) := by unfold delete exact wellAugmented_repaintBlack aug (wellAugmented_del aug x lt h)theorem storedAug_delete [Inhabited α] [DecidableEq α] (x : α) (lt : α → α → Bool) {t : AugmentedRBTree α β} (h : WellAugmented aug t) : storedAug aug (delete aug x lt t) = realAug aug (delete aug x lt t) := storedAug_eq_realAug_of_wellAugmented aug (wellAugmented_delete aug x lt h)end Genericsection Refinementvariable {β : Type} [Inhabited β] (aug : Augmentation Nat β)

Erase the cached augmentation field, projecting a Nat-keyed augmented red-black tree onto the Chapter 13 red-black tree.

def toRB : AugmentedRBTree Nat β → RBTree | empty => RBTree.empty | node c l k _ r => RBTree.node c (toRB l) k (toRB r)

The Nat strict-less-than comparison as a Bool, used to instantiate the generic insertion so that it refines Chapter 13's CLRS.­Chapter13.­RBTree.­insert.

def natLt (a b : Nat) : Bool := decide (a < b)

natLt decides strict less-than.

theorem natLt_true_iff {a b : Nat} : (natLt a b = true) ↔ a < b := by simp [natLt]

Erasing the augmentation of a mk node forgets only the cached value.

theorem toRB_mk (c : Color) (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (mk aug c l k r) = RBTree.node c (toRB l) k (toRB r) := rfl

Erasing the augmentation commutes with repainting the root black.

omit [Inhabited β] intheorem toRB_repaintBlack (t : AugmentedRBTree Nat β) : toRB (repaintBlack t) = RBTree.repaintRoot Color.black (toRB t) := by cases t with | empty => rfl | node c l k a r => rfl

Erasing the augmentation commutes with balanceLeft.

theorem toRB_balanceLeft (l : AugmentedRBTree Nat β) (y : Nat) (r : AugmentedRBTree Nat β) : toRB (balanceLeft aug l y r) = RBTree.balanceLeft (toRB l) y (toRB r) := by cases l with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the augmentation commutes with balanceRight.

theorem toRB_balanceRight (l : AugmentedRBTree Nat β) (y : Nat) (r : AugmentedRBTree Nat β) : toRB (balanceRight aug l y r) = RBTree.balanceRight (toRB l) y (toRB r) := by cases r with | empty => rfl | node c a w s b => cases c with | black => rfl | red => cases a with | empty => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl | node ca al ak as' ar => cases ca with | red => rfl | black => cases b with | empty => rfl | node cb bl bk bs br => cases cb <;> rfl

Erasing the augmentation commutes with insertFixup (at natLt).

theorem toRB_insertFixup (x : Nat) (t : AugmentedRBTree Nat β) : toRB (insertFixup aug natLt x t) = RBTree.insertFixup x (toRB t) := by induction t with | empty => rfl | node c l y a r ihl ihr => simp only [insertFixup, RBTree.insertFixup, toRB, natLt_true_iff, gt_iff_lt, apply_ite toRB, toRB_mk, toRB_balanceLeft, toRB_balanceRight, ihl, ihr]

Refinement. The augmented insertion refines the executable Chapter 13 red-black insertion: erasing the cached augmentation turns insert (at natLt) into CLRS.­Chapter13.­RBTree.­insert. This generalizes OSRBTree.toRB_insert from the size field to an arbitrary augmentation.

theorem toRB_insert (x : Nat) (t : AugmentedRBTree Nat β) : toRB (insert aug natLt x t) = RBTree.insert x (toRB t) := by unfold insert RBTree.insert rw [toRB_repaintBlack, toRB_insertFixup]

Erasure relates keys membership to Chapter 13 tree membership.

omit [Inhabited β] intheorem inTree_toRB (y : Nat) (t : AugmentedRBTree Nat β) : RBTree.InTree y (toRB t) ↔ y ∈ keys t := by induction t with | empty => simp [toRB, RBTree.InTree, keys] | node c l k a r ihl ihr => simp only [toRB, RBTree.InTree, keys, List.append_assoc, List.singleton_append, List.mem_append, List.mem_cons, ihl, ihr] tauto

Through the refinement, the Chapter 13 red-black shape invariant is maintained by the augmented insertion.

theorem redBlackShape_toRB_insert (x : Nat) {t : AugmentedRBTree Nat β} (h : RBTree.RedBlackShape (toRB t)) : RBTree.RedBlackShape (toRB (insert aug natLt x t)) := by rw [toRB_insert] exact RBTree.redBlackShape_insert h

Through the refinement, insertion preserves membership (as an inorder key).

theorem mem_keys_insert (x y : Nat) (t : AugmentedRBTree Nat β) : y ∈ keys (insert aug natLt x t) ↔ y = x ∨ y ∈ keys t := by simp only [← inTree_toRB, toRB_insert, RBTree.inTree_insert_iff]
Refinement of the deletion pipeline

Erasing the augmentation commutes with repainting the root in any colour.

omit [Inhabited β] intheorem toRB_repaintRoot (c : Color) (t : AugmentedRBTree Nat β) : toRB (repaintRoot c t) = RBTree.repaintRoot c (toRB t) := by cases t <;> rfl

Erasing the augmentation preserves the boolean black-root test.

omit [Inhabited β] intheorem rootBlack_toRB (t : AugmentedRBTree Nat β) : RBTree.rootBlack (toRB t) = rootBlack t := by cases t <;> rfl

Erasing the augmentation preserves the empty/non-empty distinction.

omit [Inhabited β] intheorem toRB_empty_iff (t : AugmentedRBTree Nat β) : toRB t = RBTree.empty ↔ t = AugmentedRBTree.empty := by cases t <;> simp [toRB]

Erasing the augmentation commutes with baldL.

theorem toRB_baldL (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (baldL aug l k r) = RBTree.baldL (toRB l) k (toRB r) := by cases l with | empty => cases r with | empty => rfl | node rc cl cy sc cr => cases rc with | black => simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases cl with | empty => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node clc cll clk cls clr => cases clc <;> simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot] | node lc a x s b => cases lc with | red => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | black => cases r with | empty => rfl | node rc cl cy sc cr => cases rc with | black => simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases cl with | empty => simp [baldL, RBTree.baldL, toRB, toRB_mk, This simp argument is unused: toRB_balanceRight Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵R̵i̵g̵h̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceRight, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node clc cll clk cls clr => cases clc <;> simp [baldL, RBTree.baldL, toRB, toRB_mk, toRB_balanceRight, toRB_repaintRoot]

Erasing the augmentation commutes with baldR.

theorem toRB_baldR (l : AugmentedRBTree Nat β) (k : Nat) (r : AugmentedRBTree Nat β) : toRB (baldR aug l k r) = RBTree.baldR (toRB l) k (toRB r) := by cases r with | empty => cases l with | empty => rfl | node lc a x s b => cases lc with | black => simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases b with | empty => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node bc bl bk bs br => cases bc <;> simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot] | node rc c y s d => cases rc with | red => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | black => cases l with | empty => rfl | node lc a x s b => cases lc with | black => simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | red => cases b with | empty => simp [baldR, RBTree.baldR, toRB, toRB_mk, This simp argument is unused: toRB_balanceLeft Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_b̵a̵l̵a̵n̵c̵e̵L̵e̵f̵t̵,̵ ̵t̵o̵R̵B̵_̵repaintRoot] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_balanceLeft, This simp argument is unused: toRB_repaintRoot Hint: Omit it from the simp argument list. simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft,̵ ̵t̵o̵R̵B̵_̵r̵e̵p̵a̵i̵n̵t̵R̵o̵o̵t̵] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_repaintRoot] | node bc bl bk bs br => cases bc <;> simp [baldR, RBTree.baldR, toRB, toRB_mk, toRB_balanceLeft, toRB_repaintRoot]

Erasing the augmentation commutes with the minimum key of splitMin.

theorem toRB_splitMin_min (t : AugmentedRBTree Nat β) : (RBTree.splitMin (toRB t)).1 = (splitMin aug t).1 := by induction t with | empty => rfl | node c l k a r ihl => cases l with | empty => rfl | node lc ll lk la lr => cases lc with | black => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack] using ihl | red => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack] using ihl

Erasing the augmentation commutes with the tree component of splitMin.

theorem toRB_splitMin_tree (t : AugmentedRBTree Nat β) : toRB (splitMin aug t).2 = (RBTree.splitMin (toRB t)).2 := by induction t with | empty => rfl | node c l k a r ihl => cases l with | empty => rfl | node lc ll lk la lr => cases lc with | black => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack, toRB_baldL] using (congrArg (fun t : RBTree => RBTree.baldL t k r.toRB) ihl) | red => simpa [splitMin, RBTree.splitMin, toRB, rootBlack, RBTree.rootBlack, toRB_mk] using (congrArg (fun t : RBTree => RBTree.node Color.red t k r.toRB) ihl)

Erasing the augmentation commutes with join.

theorem toRB_join (l r : AugmentedRBTree Nat β) : toRB (join aug l r) = RBTree.join (toRB l) (toRB r) := by cases r with | empty => cases l with | empty => rfl | node _ _ _ _ _ => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵toRB_splitMin_min, toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, This simp argument is unused: toRB_splitMin_min Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_min, This simp argument is unused: toRB_splitMin_tree Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵t̵r̵e̵e̵]̵t̲o̲R̲B̲_̲s̲p̲l̲i̲t̲M̲i̲n̲_̲m̲i̲n̲]̲ Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_tree] | node rc rl rk sr rr => cases l with | empty => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵toRB_splitMin_min, toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, This simp argument is unused: toRB_splitMin_min Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵toRB_splitMin_tree] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_min, This simp argument is unused: toRB_splitMin_tree Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, apply_ite toRB, t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵m̵i̵n̵,̵ ̵t̵o̵R̵B̵_̵s̵p̵l̵i̵t̵M̵i̵n̵_̵t̵r̵e̵e̵]̵t̲o̲R̲B̲_̲s̲p̲l̲i̲t̲M̲i̲n̲_̲m̲i̲n̲]̲ Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`toRB_splitMin_tree] | node lc ll lk sl lr => cases rc with | black => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵rootBlack, rootBlack_toRB, ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲toRB_splitMin_tree, toRB_baldR] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, rootBlack, rootBlack_toRB, toRB_splitMin_tree, toRB_baldR] rw [toRB_splitMin_min] | red => simp [join, RBTree.join, toRB_empty_iff, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [join, RBTree.join, toRB_empty_iff, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵rootBlack, rootBlack_toRB, ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲ ̲toRB_splitMin_tree, toRB_mk] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, rootBlack, rootBlack_toRB, toRB_splitMin_tree, toRB_mk] rw [toRB_splitMin_min]

Erasing the augmentation commutes with recursive deletion.

theorem toRB_del (x : Nat) (t : AugmentedRBTree Nat β) : toRB (del aug x natLt t) = RBTree.del x (toRB t) := by induction t with | empty => rfl | node c l y a r ihl ihr => by_cases hxy : x < y · have hxy' : natLt x y = true := natLt_true_iff.mpr hxy simp [del, RBTree.del, toRB, apply_ite toRB, rootBlack, hxy, hxy', rootBlack_toRB, toRB_baldL, toRB_mk, ihl] · by_cases hyx : y < x · have hxy' : natLt x y = false := by simp [natLt, hxy] have hyx' : natLt y x = true := natLt_true_iff.mpr hyx simp [del, RBTree.del, toRB, apply_ite toRB, rootBlack, hxy, hyx, hxy', hyx', rootBlack_toRB, toRB_baldR, toRB_mk, ihr] · have hxy' : natLt x y = false := by simp [natLt, hxy] have hyx' : natLt y x = false := by simp [natLt, hyx] simp [del, RBTree.del, toRB, This simp argument is unused: apply_ite toRB Hint: Omit it from the simp argument list. simp [del, RBTree.del, toRB, a̵p̵p̵l̵y̵_̵i̵t̵e̵ ̵t̵o̵R̵B̵,̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵ ̵hxy, hyx, hxy', hyx', toRB_join] Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`apply_ite toRB, hxy, hyx, hxy', hyx', toRB_join]

Deletion refinement. The augmented deletion refines the executable Chapter 13 red-black deletion: erasing the cached augmentation turns delete (at natLt) into CLRS.­Chapter13.­RBTree.­delete. This extends OSRBTree.toRB_delete from the size field to an arbitrary augmentation, closing the deletion half of the generic AugmentedRBTree refinement.

theorem toRB_delete (x : Nat) (t : AugmentedRBTree Nat β) : toRB (delete aug x natLt t) = RBTree.delete x (toRB t) := by unfold delete RBTree.delete rw [toRB_repaintBlack, toRB_del]
end Refinementsection Instances
Instance 1: the order-statistic (subtree-size) augmentation

Taking aug := sizeAug Nat recovers the order-statistic tree of §14.1: the cached field is the subtree node count, the invariant survives the executable red-black insertion, and the augmentation-erasing projection refines Chapter 13's CLRS.­Chapter13.­RBTree.­insert. This makes OSRBTree a special case of the generic interface rather than a bespoke copy.

Order-statistic instance. The subtree-size augmentation is maintained through the generic executable red-black insertion (CLRS 14.1 for size).

theorem sizeAug_wellAugmented_insert (x : Nat) {t : AugmentedRBTree Nat Nat} (h : WellAugmented (sizeAug Nat) t) : WellAugmented (sizeAug Nat) (insert (sizeAug Nat) natLt x t) := wellAugmented_insert (sizeAug Nat) natLt x h

The size augmentation's recomputed value is exactly the node count.

theorem sizeAug_realAug_eq_length (t : AugmentedRBTree Nat Nat) : realAug (sizeAug Nat) t = (keys t).length := by induction t with | empty => rfl | node c l k a r ihl ihr => simp only [realAug] rw [ihl, ihr] simp only [sizeAug, keys, List.length_append, List.length_cons, List.length_nil] omega

Order-statistic refinement. Erasing the size field turns the generic size-augmented insertion into Chapter 13's CLRS.­Chapter13.­RBTree.­insert.

theorem sizeAug_toRB_insert (x : Nat) (t : AugmentedRBTree Nat Nat) : toRB (insert (sizeAug Nat) natLt x t) = RBTree.insert x (toRB t) := toRB_insert (sizeAug Nat) x t
Instance 2: the interval-tree (maximum-high-endpoint) augmentation

Taking aug := IntervalTree.maxHighAug recovers interval trees: the cached field is the subtree's maximum high endpoint, maintained through the same generic executable insertion, with distinct intervals ordered lexicographically by low and high endpoints.

Compare intervals lexicographically by low endpoint, then high endpoint. Distinct intervals with equal low endpoints remain distinct insertion keys.

def intervalLt (i j : Interval) : Bool := decide (i.low < j.low ∨ (i.low = j.low ∧ i.high < j.high))

Interval-tree instance. The maximum-high-endpoint augmentation is maintained through the generic executable red-black insertion.

end Instancesend AugmentedRBTree

Scope and implementation notes

Imports

Current source

During the compatibility period this guide imports CLRSLean.Chapter_14. Existing declarations retain their current namespaces until the chapter-by-chapter source migration.

Coverage boundary

The third-edition Chapter 14 developments supply substantial relocated proof content, with the following fourth-edition semantic and execution interfaces:

  • §17.1 (Section_17_1_Dynamic_Order_Statistics): OS-RANK osRank/rankOf and their agreement on well-sized trees. Strict BST ordering identifies rank with the cardinality of keys below the query, including absent queries. The query descent has a logarithmic height bound.

  • §17.2 (Section_17_2_Augmenting_Data_Structures): actual cached-field insertion and rotation execution refines the legacy update on well-augmented inputs. It counts at most 5h+1 combine calls and 2h rotations. The old augmentation_update_bound is only an abstract height budget; no counted cached deletion or allocator-runtime theorem is supplied.

  • §17.3 (Section_17_3_Interval_Trees): the dynamic/static interval-tree bridge toIntervalTree/wellAugmented_toIntervalTree, full-interval lexicographic insertion preserving membership and BST order, post-insert overlap search correctness, and the Interval-keyed logarithmic query descent bound. Equal low endpoints with different high endpoints are distinct keys; only an exactly repeated interval is suppressed.

See docs/clrs-fourth-edition-map.csv for the section-level mapping and docs/migrations/clrs4.md for compatibility and deprecation policy.

CLRS, fourth edition · Chapter 17 of 35