Geb.Mathlib.Data.UnionFind.OfEdges

Imports

/-
Copyright (c) 2026 Terence Rokop. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Terence Rokop
-/
module

public import Batteries.Data.UnionFind

A size-indexed union-find and the fold over a list of edges

Batteries.UnionFind ties its size to its representation, so an index into it has type Fin self.size and every operation that changes the structure changes the index type. Sized n fixes the size as a subtype, so the indices are Fin n throughout and no cast is needed to pass one operation's index to the next. Sized.ofEdges folds Sized.union over a list of pairs, and the two theorems about it are the two directions of correctness: every listed pair is merged, and nothing beyond the listed pairs is.

Main definitions

Batteries.UnionFind.Sized — a union-find of a fixed size.

Batteries.UnionFind.Sized.discrete, Batteries.UnionFind.Sized.union, Batteries.UnionFind.Sized.root — the operations, at Fin n.

Batteries.UnionFind.Sized.ofEdges — the fold over a list of pairs.

Main statements

Batteries.UnionFind.Sized.root_ofEdges_eq_of_mem — every listed pair is merged.

Batteries.UnionFind.Sized.apply_root_ofEdges — nothing beyond the listed pairs is merged, in eliminator form.

Implementation notes

The second is stated as an eliminator — any h : Fin n → α agreeing on the listed pairs agrees on roots — rather than as a characterisation of the merged relation as the equivalence closure of the edges. The eliminator is what a coequalizer's factorisation law instantiates directly.

This module does not extract to mathlib4, Sized being a wrapper over a Batteries type, and it imports nothing outside core and Batteries. Where such content belongs is TODO.md § Upstream destination of core- and Batteries-targeted content.

Tags

union-find, disjoint set, quotient, choice-free

              @[expose] public sectionuniverse unamespace Batteries.UnionFindvariable {n : Nat}

union preserves the size.

theorem size_union (self : UnionFind) (x y : Fin self.size) :
    (self.union x y).size = self.size := byself:UnionFindx:Fin self.sizey:Fin self.size⊢ (self.union x y).size = self.size
  unfold unionself:UnionFindx:Fin self.sizey:Fin self.size⊢ (match self.find x with
    | ⟨self₁, ⟨rx, ex⟩⟩ =>
      have hy := ⋯;
      match eq : self₁.find ⟨↑y, hy⟩ with
      | ⟨self₂, ⟨ry, ey⟩⟩ => self₂.link ⟨↑rx, ⋯⟩ ry ⋯).size =
  self.size; simp [UnionFind.size]All goals completed! 🐙

push adds one to the size.

theorem size_push (self : UnionFind) : self.push.size = self.size + 1 := byself:UnionFind⊢ self.push.size = self.size + 1
  unfold pushself:UnionFind⊢ { arr := self.arr.push { parent := self.arr.size, rank := 0 }, parentD_lt := ⋯, rankD_lt := ⋯ }.size = self.size + 1; simp [UnionFind.size]All goals completed! 🐙

A union-find whose size is fixed, so that its indices are Fin n and no operation changes their type.

def Sized (n : Nat) : Type := {u : UnionFind // u.size = n}

The discrete partition on n elements: n pushes onto the empty structure.

def Sized.discrete (n : Nat) : Sized n :=
  Nat.rec (motive := fun m ↦ Sized m) ⟨.empty, rfl⟩
    (fun _ v ↦ ⟨v.1.push, byn✝:Natn:Natx✝:Natv:Sized x✝⊢ v.val.push.size = x✝.succ rw [size_push,n✝:Natn:Natx✝:Natv:Sized x✝⊢ v.val.size + 1 = x✝.succ v.2n✝:Natn:Natx✝:Natv:Sized x✝⊢ x✝ + 1 = x✝.succ]All goals completed! 🐙⟩) n

Merge the classes of two indices.

def Sized.union (v : Sized n) (x y : Fin n) : Sized n :=
  ⟨v.1.unionN x y v.2.symm, byn:Natv:Sized nx:Fin ny:Fin n⊢ (v.val.unionN x y ⋯).size = n obtain ⟨u, rfl⟩ := vu:UnionFindx:Fin u.sizey:Fin u.size⊢ (⟨u, ⋯⟩.val.unionN x y ⋯).size = u.size; exact size_union u x yAll goals completed! 🐙⟩

The representative of an index's class, as an index.

def Sized.root (v : Sized n) (x : Fin n) : Fin n :=
  ⟨v.1.rootD x, byn:Natv:Sized nx:Fin n⊢ v.val.rootD ↑x < n obtain ⟨u, rfl⟩ := vu:UnionFindx:Fin u.size⊢ ⟨u, ⋯⟩.val.rootD ↑x < u.size; exact UnionFind.rootD_lt.mpr x.isLtAll goals completed! 🐙⟩

The union-find obtained by merging every listed pair.

def Sized.ofEdges (n : Nat) (l : List (Fin n × Fin n)) : Sized n :=
  l.foldl (fun v p ↦ v.union p.1 p.2) (discrete n)

Two indices have the same root exactly when they are equivalent.

theorem Sized.root_eq_iff {v : Sized n} {a b : Fin n} :
    v.root a = v.root b ↔ v.1.Equiv a b := Fin.ext_iff

Batteries.UnionFind.equiv_union restated at Sized.union. The Nat arguments match Batteries' Equiv; the Fin n arguments the other lemmas pass are coerced.

theorem Sized.equiv_union {v : Sized n} {x y : Fin n} {a b : Nat} :
    (v.union x y).1.Equiv a b ↔
      v.1.Equiv a b ∨ v.1.Equiv a x ∧ v.1.Equiv y b
                    ∨ v.1.Equiv a y ∧ v.1.Equiv x b := byn:Natv:Sized nx:Fin ny:Fin na:Natb:Nat⊢ (v.union x y).val.Equiv a b ↔
  v.val.Equiv a b ∨ v.val.Equiv a ↑x ∧ v.val.Equiv (↑y) b ∨ v.val.Equiv a ↑y ∧ v.val.Equiv (↑x) b
  obtain ⟨u, rfl⟩ := va:Natb:Natu:UnionFindx:Fin u.sizey:Fin u.size⊢ (union ⟨u, ⋯⟩ x y).val.Equiv a b ↔
  ⟨u, ⋯⟩.val.Equiv a b ∨
    ⟨u, ⋯⟩.val.Equiv a ↑x ∧ ⟨u, ⋯⟩.val.Equiv (↑y) b ∨ ⟨u, ⋯⟩.val.Equiv a ↑y ∧ ⟨u, ⋯⟩.val.Equiv (↑x) b
  exact UnionFind.equiv_unionAll goals completed! 🐙

Every index is its own root in the discrete partition.

theorem Sized.rootD_discrete (m x : Nat) : (discrete m).1.rootD x = x :=
  Nat.rec (motive := fun k ↦ (discrete k).1.rootD x = x)
    UnionFind.rootD_empty (fun _ ih ↦ (UnionFind.root_push).trans ih) m

Sized.rootD_discrete at Fin n.

theorem Sized.root_discrete (x : Fin n) : (discrete n).root x = x :=
  Fin.ext (rootD_discrete n x)

A root is its own root.

theorem Sized.root_root (v : Sized n) (x : Fin n) :
    v.root (v.root x) = v.root x := Fin.ext UnionFind.rootD_rootD

Equivalence in an accumulator survives the fold.

theorem Sized.equiv_foldl_of_equiv (l : List (Fin n × Fin n))
    (a b : Fin n) (v : Sized n) (hv : v.1.Equiv a b) :
    (l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).1.Equiv
      a b :=
  List.rec (motive := fun (l : List (Fin n × Fin n)) ↦ ∀ (v : Sized n), v.1.Equiv a b →
      (l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).1.Equiv a b)
    (fun _ hv ↦ hv)
    (fun p _ ih v hv ↦ ih (v.union p.1 p.2) (Sized.equiv_union.mpr (Or.inl hv)))
    l v hv

A listed pair is equivalent after the fold, from any accumulator.

theorem Sized.equiv_foldl_of_mem (l : List (Fin n × Fin n))
    (a b : Fin n) (hab : (a, b) ∈ l) (v : Sized n) :
    (l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).1.Equiv
      a b :=
  List.rec (motive := fun (l : List (Fin n × Fin n)) ↦ (a, b) ∈ l → ∀ (v : Sized n),
      (l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).1.Equiv a b)
    (fun hab ↦ absurd hab List.not_mem_nil)
    (fun p _ ih hab v ↦ byn:Natl:List (Fin n × Fin n)a:Fin nb:Fin nhab✝:(a, b) ∈ lv✝:Sized np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(a, b) ∈ x✝ → ∀ (v : Sized n), (List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).val.Equiv ↑a ↑bhab:(a, b) ∈ p :: x✝v:Sized n⊢ (List.foldl (fun v p ↦ v.union p.fst p.snd) v (p :: x✝)).val.Equiv ↑a ↑b
      cases List.mem_cons.mp hab with
      | inl hp =>inln:Natl:List (Fin n × Fin n)a:Fin nb:Fin nhab✝:(a, b) ∈ lv✝:Sized np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(a, b) ∈ x✝ → ∀ (v : Sized n), (List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).val.Equiv ↑a ↑bhab:(a, b) ∈ p :: x✝v:Sized nhp:(a, b) = p⊢ (List.foldl (fun v p ↦ v.union p.fst p.snd) v (p :: x✝)).val.Equiv ↑a ↑b
        subst hpinln:Natl:List (Fin n × Fin n)a:Fin nb:Fin nhab✝:(a, b) ∈ lv✝:Sized nx✝:List (Fin n × Fin n)ih:(a, b) ∈ x✝ → ∀ (v : Sized n), (List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).val.Equiv ↑a ↑bv:Sized nhab:(a, b) ∈ (a, b) :: x✝⊢ (List.foldl (fun v p ↦ v.union p.fst p.snd) v ((a, b) :: x✝)).val.Equiv ↑a ↑b
        exact equiv_foldl_of_equiv _ a b (v.union a b)
          (Sized.equiv_union.mpr (Or.inr (Or.inl ⟨rfl, rfl⟩)))All goals completed! 🐙
      | inr ht =>inrn:Natl:List (Fin n × Fin n)a:Fin nb:Fin nhab✝:(a, b) ∈ lv✝:Sized np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(a, b) ∈ x✝ → ∀ (v : Sized n), (List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).val.Equiv ↑a ↑bhab:(a, b) ∈ p :: x✝v:Sized nht:(a, b) ∈ x✝⊢ (List.foldl (fun v p ↦ v.union p.fst p.snd) v (p :: x✝)).val.Equiv ↑a ↑b exact ih ht (v.union p.1 p.2)All goals completed! 🐙)
    l hab v

A function agreeing on every listed pair, and on the accumulator's roots, agrees on the roots after the fold.

theorem Sized.apply_root_foldl {α : Type u} {h : Fin n → α}
    (l : List (Fin n × Fin n)) (hl : ∀ p ∈ l, h p.1 = h p.2)
    (v : Sized n) (hv : ∀ x, h (v.root x) = h x) (x : Fin n) :
    h ((l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).root
      x) = h x :=
  List.rec (motive := fun (l : List (Fin n × Fin n)) ↦ (∀ p ∈ l, h p.1 = h p.2) →
      ∀ (v : Sized n), (∀ z, h (v.root z) = h z) →
      h ((l.foldl (fun (v : Sized n) (p : Fin n × Fin n) ↦ v.union p.1 p.2) v).root x) = h x)
    (fun _ _ hv ↦ hv x)
    (fun p _ ih hl v hv ↦ ih (fun q hq ↦ hl q (List.mem_cons_of_mem p hq))
      (v.union p.1 p.2) (fun z ↦ byn:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin n⊢ h ((v.union p.fst p.snd).root z) = h z
        have key : ∀ c d : Fin n, v.1.Equiv c d → h c = h d := fun c d hcd ↦
          ((hv c).symm.trans (congrArg h (root_eq_iff.mpr hcd))).trans (hv d)n:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h d⊢ h ((v.union p.fst p.snd).root z) = h z
        have hz : (v.union p.1 p.2).1.Equiv z ((v.union p.1 p.2).root z) :=
          UnionFind.rootD_rootD.symmn:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h dhz:(v.union p.fst p.snd).val.Equiv ↑z ↑((v.union p.fst p.snd).root z)⊢ h ((v.union p.fst p.snd).root z) = h z
        cases Sized.equiv_union.mp hz with
        | inl hsame =>inln:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h dhz:(v.union p.fst p.snd).val.Equiv ↑z ↑((v.union p.fst p.snd).root z)hsame:v.val.Equiv ↑z ↑((v.union p.fst p.snd).root z)⊢ h ((v.union p.fst p.snd).root z) = h z exact (key _ _ hsame).symmAll goals completed! 🐙
        | inr hcross =>inrn:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h dhz:(v.union p.fst p.snd).val.Equiv ↑z ↑((v.union p.fst p.snd).root z)hcross:v.val.Equiv ↑z ↑p.fst ∧ v.val.Equiv ↑p.snd ↑((v.union p.fst p.snd).root z) ∨
  v.val.Equiv ↑z ↑p.snd ∧ v.val.Equiv ↑p.fst ↑((v.union p.fst p.snd).root z)⊢ h ((v.union p.fst p.snd).root z) = h z
          cases hcross with
          | inl hfwd =>inr.inln:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h dhz:(v.union p.fst p.snd).val.Equiv ↑z ↑((v.union p.fst p.snd).root z)hfwd:v.val.Equiv ↑z ↑p.fst ∧ v.val.Equiv ↑p.snd ↑((v.union p.fst p.snd).root z)⊢ h ((v.union p.fst p.snd).root z) = h z exact ((key _ _ hfwd.1).trans
              ((hl p List.mem_cons_self).trans (key _ _ hfwd.2))).symmAll goals completed! 🐙
          | inr hbwd =>inr.inrn:Natα:Type uh:Fin n → αl:List (Fin n × Fin n)hl✝:∀ (p : Fin n × Fin n), p ∈ l → h p.fst = h p.sndv✝:Sized nhv✝:∀ (x : Fin n), h (v✝.root x) = h xx:Fin np:Fin n × Fin nx✝:List (Fin n × Fin n)ih:(∀ (p : Fin n × Fin n), p ∈ x✝ → h p.fst = h p.snd) →
  ∀ (v : Sized n),
    (∀ (z : Fin n), h (v.root z) = h z) → h ((List.foldl (fun v p ↦ v.union p.fst p.snd) v x✝).root x) = h xhl:∀ (p_1 : Fin n × Fin n), p_1 ∈ p :: x✝ → h p_1.fst = h p_1.sndv:Sized nhv:∀ (z : Fin n), h (v.root z) = h zz:Fin nkey:∀ (c d : Fin n), v.val.Equiv ↑c ↑d → h c = h dhz:(v.union p.fst p.snd).val.Equiv ↑z ↑((v.union p.fst p.snd).root z)hbwd:v.val.Equiv ↑z ↑p.snd ∧ v.val.Equiv ↑p.fst ↑((v.union p.fst p.snd).root z)⊢ h ((v.union p.fst p.snd).root z) = h z exact ((key _ _ hbwd.1).trans
              ((hl p List.mem_cons_self).symm.trans (key _ _ hbwd.2))).symmAll goals completed! 🐙))
    l hl v hv

Every listed pair is merged.

theorem Sized.root_ofEdges_eq_of_mem {l : List (Fin n × Fin n)}
    {a b : Fin n} (hab : (a, b) ∈ l) :
    (ofEdges n l).root a = (ofEdges n l).root b :=
  root_eq_iff.mpr (equiv_foldl_of_mem l a b hab (discrete n))

Nothing beyond the listed pairs is merged: a function agreeing on every listed pair agrees on roots.

theorem Sized.apply_root_ofEdges {α : Type u} {l : List (Fin n × Fin n)}
    {h : Fin n → α} (hl : ∀ p ∈ l, h p.1 = h p.2) (x : Fin n) :
    h ((ofEdges n l).root x) = h x :=
  apply_root_foldl l hl (discrete n) (fun z ↦ congrArg h (root_discrete z)) x

end Batteries.UnionFind