feat(UnionFind): union by rank with verified O(log n) find

Cslib.lean

@@ -2,6 +2,10 @@ module  -- shake: keep-all --deprecated_module: ignore
 
 public import Cslib.Algorithms.Lean.MergeSort.MergeSort
 public import Cslib.Algorithms.Lean.TimeM
+public import Cslib.Algorithms.Lean.UnionFind.Basic
+public import Cslib.Algorithms.Lean.UnionFind.Cost
+public import Cslib.Algorithms.Lean.UnionFind.Operations
+public import Cslib.Algorithms.Lean.UnionFind.RankBound
 public import Cslib.Computability.Automata.Acceptors.Acceptor
 public import Cslib.Computability.Automata.Acceptors.OmegaAcceptor
 public import Cslib.Computability.Automata.DA.Basic

Cslib/Algorithms/Lean/UnionFind/Basic.lean

@@ -0,0 +1,102 @@
+/-
+Copyright (c) 2026 Aviv Bar Natan. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Aviv Bar Natan
+-/
+
+module
+
+public import Cslib.Init
+public import Mathlib.Data.Finset.Card
+public import Mathlib.Data.Fintype.Basic
+
+/-!
+# Union-Find Forest Model
+
+This file defines a mathematical model of a union-by-rank union-find forest over `Fin n`.
+
+## Main definitions
+
+* `UnionFind`: A forest over `n` elements given by parent and rank functions.
+* `UnionFind.IsRoot`: A node is a root iff it is its own parent.
+* `UnionFind.reaches`: Reachability by following parent pointers.
+* `UnionFind.subtree`: The set of nodes that reach a given root.
+* `UnionFind.subtreeSize`: The cardinality of a subtree.
+* `UnionFind.WellFormed`: Well-formedness for union by rank.
+
+## Main results
+
+* `UnionFind.self_mem_subtree`: Every node belongs to its own subtree.
+* `UnionFind.one_le_subtreeSize`: Every subtree is nonempty.
+* `UnionFind.subtreeSize_le`: Every subtree has at most `n` nodes.
+-/
+
+@[expose] public section
+
+set_option autoImplicit false
+
+namespace Cslib.Algorithms.Lean
+
+variable {n : ℕ}
+
+/-- A union-find forest over `n` elements, given by parent and rank functions. -/
+structure UnionFind (n : ℕ) where
+  /-- Parent pointer; roots point to themselves. -/
+  parent : Fin n → Fin n
+  /-- Rank (an upper bound on subtree height) of each node. -/
+  rank : Fin n → ℕ
+
+namespace UnionFind
+
+variable (uf : UnionFind n)
+
+/-- `i` is a root iff it is its own parent. -/
+def IsRoot (i : Fin n) : Prop := uf.parent i = i
+
+instance (i : Fin n) : Decidable (uf.IsRoot i) := by unfold IsRoot; infer_instance
+
+/-- `j` reaches `i` by following parent pointers some number of times. -/
+def reaches (j i : Fin n) : Prop := ∃ k, uf.parent^[k] j = i
+
+open Classical in
+/-- The set of all nodes `j` that reach `i` by following parent pointers. -/
+noncomputable def subtree (i : Fin n) : Finset (Fin n) :=
+  Finset.univ.filter (fun j => uf.reaches j i)
+
+/-- The number of nodes in the subtree rooted at `i`. -/
+noncomputable def subtreeSize (i : Fin n) : ℕ := (uf.subtree i).card
+
+/-- Well-formedness for union by rank: ranks increase toward roots, and a rank-`r` node
+owns at least `2 ^ r` nodes. -/
+structure WellFormed : Prop where
+  /-- Non-root nodes have strictly smaller rank than their parent. -/
+  rank_lt : ∀ i, ¬ uf.IsRoot i → uf.rank i < uf.rank (uf.parent i)
+  /-- A rank-`r` node owns at least `2 ^ r` nodes. -/
+  size_ge : ∀ i, 2 ^ uf.rank i ≤ uf.subtreeSize i
+
+/-- Every node reaches itself in zero steps. -/
+@[simp] lemma reaches_self (i : Fin n) : uf.reaches i i := ⟨0, rfl⟩
+
+open Classical in
+/-- A node is always in its own subtree. -/
+@[simp] lemma self_mem_subtree (i : Fin n) : i ∈ uf.subtree i := by
+  simp only [subtree, Finset.mem_filter, Finset.mem_univ, reaches_self, and_self]
+
+/-- Every subtree is nonempty. -/
+@[simp] lemma one_le_subtreeSize (i : Fin n) : 1 ≤ uf.subtreeSize i := by
+  unfold subtreeSize
+  exact Finset.card_pos.mpr ⟨i, uf.self_mem_subtree i⟩
+
+open Classical in
+/-- A subtree has at most `n` nodes. -/
+lemma subtreeSize_le (i : Fin n) : uf.subtreeSize i ≤ n := by
+  unfold subtreeSize subtree
+  have hle : (Finset.univ.filter (fun j => uf.reaches j i)).card ≤ Finset.univ.card :=
+    Finset.card_le_card (Finset.filter_subset _ _)
+  have huniv : (Finset.univ : Finset (Fin n)).card = n := by
+    simp [Finset.card_def, Fin.univ_def, List.length_finRange]
+  omega
+
+end UnionFind
+
+end Cslib.Algorithms.Lean

Cslib/Algorithms/Lean/UnionFind/Cost.lean

@@ -0,0 +1,81 @@
+/-
+Copyright (c) 2026 Aviv Bar Natan. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Aviv Bar Natan
+-/
+
+module
+
+public import Cslib.Init
+public import Cslib.Algorithms.Lean.UnionFind.RankBound
+public import Cslib.Algorithms.Lean.TimeM
+
+/-!
+# Timed `find` for Union-Find
+
+This file expresses the `find` operation of union-find in the `TimeM` cost monad,
+charging one tick per parent dereference. We prove that `find` returns the correct
+root (`find_ret`), that its tick cost equals the find depth (`find_time`), and that
+in a well-formed union-by-rank forest it runs in at most `⌊log₂ n⌋` ticks
+(`find_time_le_log`).
+
+Following the `TimeM` convention, correctness is proved on `.ret` and complexity on
+`.time`.
+
+## Main definitions
+
+* `UnionFind.findAux`: follow `k` parent pointers from `x`, charging one tick per hop.
+* `UnionFind.find`: follow parent pointers to the root, charging one tick per dereference.
+
+## Main results
+
+* `UnionFind.find_ret`: `find` returns the root of `i`.
+* `UnionFind.find_time`: `find`'s tick cost equals the find depth.
+* `UnionFind.find_time_le_log`: `find` runs in at most `⌊log₂ n⌋` ticks.
+-/
+
+@[expose] public section
+
+set_option autoImplicit false
+
+namespace Cslib.Algorithms.Lean.UnionFind
+
+open Cslib.Algorithms.Lean (TimeM)
+
+variable {n : ℕ} (uf : UnionFind n)
+
+/-- Auxiliary for `find`: follow `k` parent pointers from `x`, charging one tick per hop. -/
+def findAux (uf : UnionFind n) : ℕ → Fin n → TimeM ℕ (Fin n)
+  | 0, x => pure x
+  | (k + 1), x => do ✓ uf.findAux k (uf.parent x)
+
+/-- The value returned by `findAux k x` is the `k`-th parent iterate of `x`. -/
+@[simp] lemma findAux_ret (k : ℕ) (x : Fin n) : (uf.findAux k x).ret = uf.parent^[k] x := by
+  induction k generalizing x with
+  | zero => simp [findAux]
+  | succ k ih => simp [findAux, ih]
+
+/-- `findAux k x` costs exactly `k` ticks. -/
+@[simp] lemma findAux_time (k : ℕ) (x : Fin n) : (uf.findAux k x).time = k := by
+  induction k generalizing x with
+  | zero => simp [findAux]
+  | succ k ih => simp [findAux, ih]; omega
+
+/-- `find i` follows parent pointers to `i`'s root, charging one tick per dereference. -/
+noncomputable def find (hwf : uf.WellFormed) (i : Fin n) : TimeM ℕ (Fin n) :=
+  uf.findAux (uf.findDepth hwf i) i
+
+/-- `find` returns the root of `i`. -/
+lemma find_ret (hwf : uf.WellFormed) (i : Fin n) : (uf.find hwf i).ret = uf.findRoot hwf i := by
+  simp [find, findAux_ret, findRoot]
+
+/-- `find`'s tick cost is exactly the find depth. -/
+lemma find_time (hwf : uf.WellFormed) (i : Fin n) : (uf.find hwf i).time = uf.findDepth hwf i := by
+  simp [find, findAux_time]
+
+/-- `find` runs in at most `⌊log₂ n⌋` ticks. -/
+theorem find_time_le_log (hwf : uf.WellFormed) (i : Fin n) :
+    (uf.find hwf i).time ≤ Nat.log 2 n := by
+  rw [uf.find_time hwf i]; exact uf.findDepth_le_log hwf i
+
+end Cslib.Algorithms.Lean.UnionFind