structure Batteries.UFNode : Type

Union-find node type

• parent : Nat

  Parent of node

• rank : Nat

  Rank of node

def Batteries.UnionFind.panicWith {α : Type u_1} (v : α) (msg : String) : α

Panic with return value

Equations

• Batteries.UnionFind.panicWith v msg = panic msg

@[simp]

theorem Batteries.UnionFind.panicWith_eq {α : Type u_1} (v : α) (msg : String) : Batteries.UnionFind.panicWith v msg = v

def Batteries.UnionFind.parentD (arr : Array Batteries.UFNode) (i : Nat) : Nat

Parent of a union-find node, defaults to self when the node is a root

Equations

• Batteries.UnionFind.parentD arr i = if h : i < arr.size then (arr.get ⟨i, h⟩).parent else i

def Batteries.UnionFind.rankD (arr : Array Batteries.UFNode) (i : Nat) : Nat

Rank of a union-find node, defaults to 0 when the node is a root

Equations

• Batteries.UnionFind.rankD arr i = if h : i < arr.size then (arr.get ⟨i, h⟩).rank else 0

theorem Batteries.UnionFind.parentD_eq {arr : Array Batteries.UFNode} {i : Fin arr.size} : Batteries.UnionFind.parentD arr ↑i = (arr.get i).parent

theorem Batteries.UnionFind.parentD_eq' {arr : Array Batteries.UFNode} {i : Nat} (h : i < arr.size) : Batteries.UnionFind.parentD arr i = (arr.get ⟨i, h⟩).parent

theorem Batteries.UnionFind.rankD_eq {arr : Array Batteries.UFNode} {i : Fin arr.size} : Batteries.UnionFind.rankD arr ↑i = (arr.get i).rank

theorem Batteries.UnionFind.rankD_eq' {arr : Array Batteries.UFNode} {i : Nat} (h : i < arr.size) : Batteries.UnionFind.rankD arr i = (arr.get ⟨i, h⟩).rank

theorem Batteries.UnionFind.parentD_of_not_lt {i : Nat} {arr : Array Batteries.UFNode} : ¬i < arr.size → Batteries.UnionFind.parentD arr i = i

theorem Batteries.UnionFind.lt_of_parentD {arr : Array Batteries.UFNode} {i : Nat} : Batteries.UnionFind.parentD arr i ≠ i → i < arr.size

theorem Batteries.UnionFind.parentD_set {arr : Array Batteries.UFNode} {x : Fin arr.size} {v : Batteries.UFNode} {i : Nat} : Batteries.UnionFind.parentD (arr.set x v) i = if ↑x = i then v.parent else Batteries.UnionFind.parentD arr i

theorem Batteries.UnionFind.rankD_set {arr : Array Batteries.UFNode} {x : Fin arr.size} {v : Batteries.UFNode} {i : Nat} : Batteries.UnionFind.rankD (arr.set x v) i = if ↑x = i then v.rank else Batteries.UnionFind.rankD arr i

structure Batteries.UnionFind : Type

Union-find data structure

The UnionFind structure is an implementation of disjoint-set data structure that uses path compression to make the primary operations run in amortized nearly linear time. The nodes of a UnionFind structure s are natural numbers smaller than s.size. The structure associates with a canonical representative from its equivalence class. The structure can be extended using the push operation and equivalence classes can be updated using the union operation.

The main operations for UnionFind are:

• empty/mkEmpty are used to create a new empty structure.
• size returns the size of the data structure.
• push adds a new node to a structure, unlinked to any other node.
• union links two nodes of the data structure, joining their equivalence classes, and performs path compression.
• find returns the canonical representative of a node and updates the data structure using path compression.
• root returns the canonical representative of a node without altering the data structure.
• checkEquiv checks whether two nodes have the same canonical representative and updates the structure using path compression.

Most use cases should prefer find over root to benefit from the speedup from path-compression.

The main operations use Fin s.size to represent nodes of the union-find structure. Some alternatives are provided:

• unionN, findN, rootN, checkEquivN use Fin n with a proof that n = s.size.
• union!, find!, root!, checkEquiv! use Nat and panic when the indices are out of bounds.
• findD, rootD, checkEquivD use Nat and treat out of bound indices as isolated nodes.

The noncomputable relation UnionFind.Equiv is provided to use the equivalence relation from a UnionFind structure in the context of proofs.

• arr : Array Batteries.UFNode

  Array of union-find nodes

• parentD_lt : ∀ {i : Nat}, i < self.arr.size → Batteries.UnionFind.parentD self.arr i < self.arr.size

  Validity for parent nodes

• rankD_lt : ∀ {i : Nat}, Batteries.UnionFind.parentD self.arr i ≠ i → Batteries.UnionFind.rankD self.arr i < Batteries.UnionFind.rankD self.arr (Batteries.UnionFind.parentD self.arr i)

  Validity for rank

theorem Batteries.UnionFind.parentD_lt (self : Batteries.UnionFind) {i : Nat} : i < self.arr.size → Batteries.UnionFind.parentD self.arr i < self.arr.size

Validity for parent nodes

theorem Batteries.UnionFind.rankD_lt (self : Batteries.UnionFind) {i : Nat} : Batteries.UnionFind.parentD self.arr i ≠ i → Batteries.UnionFind.rankD self.arr i < Batteries.UnionFind.rankD self.arr (Batteries.UnionFind.parentD self.arr i)

Validity for rank

@[reducible, inline]

abbrev Batteries.UnionFind.size (self : Batteries.UnionFind) : Nat

Size of union-find structure.

Equations

• self.size = self.arr.size

def Batteries.UnionFind.mkEmpty (c : Nat) : Batteries.UnionFind

Create an empty union-find structure with specific capacity

Equations

• Batteries.UnionFind.mkEmpty c = { arr := Array.mkEmpty c, parentD_lt := ⋯, rankD_lt := ⋯ }

def Batteries.UnionFind.empty : Batteries.UnionFind

Empty union-find structure

Equations

• Batteries.UnionFind.empty = Batteries.UnionFind.mkEmpty 0

instance Batteries.UnionFind.instEmptyCollection : EmptyCollection Batteries.UnionFind

Equations

• Batteries.UnionFind.instEmptyCollection = { emptyCollection := Batteries.UnionFind.empty }

@[reducible, inline]

abbrev Batteries.UnionFind.parent (self : Batteries.UnionFind) (i : Nat) : Nat

Parent of union-find node

Equations

• self.parent i = Batteries.UnionFind.parentD self.arr i

theorem Batteries.UnionFind.parent'_lt (self : Batteries.UnionFind) (i : Fin self.size) : (self.arr.get i).parent < self.size

theorem Batteries.UnionFind.parent_lt (self : Batteries.UnionFind) (i : Nat) : self.parent i < self.size ↔ i < self.size

@[reducible, inline]

abbrev Batteries.UnionFind.rank (self : Batteries.UnionFind) (i : Nat) : Nat

Rank of union-find node

Equations

• self.rank i = Batteries.UnionFind.rankD self.arr i

theorem Batteries.UnionFind.rank_lt {self : Batteries.UnionFind} {i : Nat} : self.parent i ≠ i → self.rank i < self.rank (self.parent i)

theorem Batteries.UnionFind.rank'_lt (self : Batteries.UnionFind) (i : Fin self.size) : (self.arr.get i).parent ≠ ↑i → self.rank ↑i < self.rank (self.arr.get i).parent

noncomputable def Batteries.UnionFind.rankMax (self : Batteries.UnionFind) : Nat

Maximum rank of nodes in a union-find structure

Equations

• self.rankMax = Array.foldr (fun (x : Batteries.UFNode) => max x.rank) 0 self.arr self.arr.size + 1

theorem Batteries.UnionFind.rank'_lt_rankMax (self : Batteries.UnionFind) (i : Fin self.size) : (self.arr.get i).rank < self.rankMax

theorem Batteries.UnionFind.rank'_lt_rankMax.go {l : List Batteries.UFNode} {x : Batteries.UFNode} : x ∈ l → x.rank ≤ List.foldr (fun (x : Batteries.UFNode) => max x.rank) 0 l

theorem Batteries.UnionFind.rankD_lt_rankMax (self : Batteries.UnionFind) (i : Nat) : Batteries.UnionFind.rankD self.arr i < self.rankMax

theorem Batteries.UnionFind.lt_rankMax (self : Batteries.UnionFind) (i : Nat) : self.rank i < self.rankMax

theorem Batteries.UnionFind.push_rankD {i : Nat} (arr : Array Batteries.UFNode) : Batteries.UnionFind.rankD (arr.push { parent := arr.size, rank := 0 }) i = Batteries.UnionFind.rankD arr i

theorem Batteries.UnionFind.push_parentD {i : Nat} (arr : Array Batteries.UFNode) : Batteries.UnionFind.parentD (arr.push { parent := arr.size, rank := 0 }) i = Batteries.UnionFind.parentD arr i

def Batteries.UnionFind.push (self : Batteries.UnionFind) : Batteries.UnionFind

Add a new node to a union-find structure, unlinked with any other nodes

Equations

• self.push = { arr := self.arr.push { parent := self.arr.size, rank := 0 }, parentD_lt := ⋯, rankD_lt := ⋯ }

def Batteries.UnionFind.root (self : Batteries.UnionFind) (x : Fin self.size) : Fin self.size

Root of a union-find node.

Equations

• self.root x = if h : (self.arr.get x).parent = ↑x then x else let_fun this := ⋯; self.root ⟨(self.arr.get x).parent, ⋯⟩

def Batteries.UnionFind.rootN {n : Nat} (self : Batteries.UnionFind) (x : Fin n) (h : n = self.size) : Fin n

Root of a union-find node.

Equations

• self.rootN x h = match n, h, x with | .(self.size), ⋯, x => self.root x

def Batteries.UnionFind.root! (self : Batteries.UnionFind) (x : Nat) : Nat

Root of a union-find node. Panics if index is out of bounds.

Equations

• self.root! x = if h : x < self.size then ↑(self.root ⟨x, h⟩) else Batteries.UnionFind.panicWith x "index out of bounds"

def Batteries.UnionFind.rootD (self : Batteries.UnionFind) (x : Nat) : Nat

Root of a union-find node. Returns input if index is out of bounds.

Equations

• self.rootD x = if h : x < self.size then ↑(self.root ⟨x, h⟩) else x

theorem Batteries.UnionFind.parent_root (self : Batteries.UnionFind) (x : Fin self.size) : (self.arr.get (self.root x)).parent = ↑(self.root x)

theorem Batteries.UnionFind.parent_rootD (self : Batteries.UnionFind) (x : Nat) : self.parent (self.rootD x) = self.rootD x

theorem Batteries.UnionFind.rootD_parent (self : Batteries.UnionFind) (x : Nat) : self.rootD (self.parent x) = self.rootD x

theorem Batteries.UnionFind.rootD_lt {self : Batteries.UnionFind} {x : Nat} : self.rootD x < self.size ↔ x < self.size

theorem Batteries.UnionFind.rootD_eq_self {self : Batteries.UnionFind} {x : Nat} : self.rootD x = x ↔ self.parent x = x

theorem Batteries.UnionFind.rootD_rootD {self : Batteries.UnionFind} {x : Nat} : self.rootD (self.rootD x) = self.rootD x

theorem Batteries.UnionFind.rootD_ext {m1 : Batteries.UnionFind} {m2 : Batteries.UnionFind} (H : ∀ (x : Nat), m1.parent x = m2.parent x) {x : Nat} : m1.rootD x = m2.rootD x

theorem Batteries.UnionFind.le_rank_root {self : Batteries.UnionFind} {x : Nat} : self.rank x ≤ self.rank (self.rootD x)

theorem Batteries.UnionFind.lt_rank_root {self : Batteries.UnionFind} {x : Nat} : self.rank x < self.rank (self.rootD x) ↔ self.parent x ≠ x

structure Batteries.UnionFind.FindAux (n : Nat) : Type

Auxiliary data structure for find operation

• s : Array Batteries.UFNode

  Array of nodes

• root : Fin n

  Index of root node

• size_eq : self.s.size = n

  Size requirement

theorem Batteries.UnionFind.FindAux.size_eq {n : Nat} (self : Batteries.UnionFind.FindAux n) : self.s.size = n

Size requirement

def Batteries.UnionFind.findAux (self : Batteries.UnionFind) (x : Fin self.size) : Batteries.UnionFind.FindAux self.size

Auxiliary function for find operation

Equations

• One or more equations did not get rendered due to their size.

theorem Batteries.UnionFind.findAux_root {self : Batteries.UnionFind} {x : Fin self.size} : (self.findAux x).root = self.root x

theorem Batteries.UnionFind.findAux_s {self : Batteries.UnionFind} {x : Fin self.size} : (self.findAux x).s = if (self.arr.get x).parent = ↑x then self.arr else (self.findAux ⟨(self.arr.get x).parent, ⋯⟩).s.modify ↑x fun (s : Batteries.UFNode) => { parent := self.rootD ↑x, rank := s.rank }

theorem Batteries.UnionFind.rankD_findAux {i : Nat} {self : Batteries.UnionFind} {x : Fin self.size} : Batteries.UnionFind.rankD (self.findAux x).s i = self.rank i

theorem Batteries.UnionFind.parentD_findAux {i : Nat} {self : Batteries.UnionFind} {x : Fin self.size} : Batteries.UnionFind.parentD (self.findAux x).s i = if i = ↑x then self.rootD ↑x else Batteries.UnionFind.parentD (self.findAux ⟨(self.arr.get x).parent, ⋯⟩).s i

theorem Batteries.UnionFind.parentD_findAux_rootD {self : Batteries.UnionFind} {x : Fin self.size} : Batteries.UnionFind.parentD (self.findAux x).s (self.rootD ↑x) = self.rootD ↑x

theorem Batteries.UnionFind.parentD_findAux_lt {i : Nat} {self : Batteries.UnionFind} {x : Fin self.size} (h : i < self.size) : Batteries.UnionFind.parentD (self.findAux x).s i < self.size

theorem Batteries.UnionFind.parentD_findAux_or (self : Batteries.UnionFind) (x : Fin self.size) (i : Nat) : Batteries.UnionFind.parentD (self.findAux x).s i = self.rootD i ∧ self.rootD i = self.rootD ↑x ∨ Batteries.UnionFind.parentD (self.findAux x).s i = self.parent i

theorem Batteries.UnionFind.lt_rankD_findAux {i : Nat} {self : Batteries.UnionFind} {x : Fin self.size} : Batteries.UnionFind.parentD (self.findAux x).s i ≠ i → self.rank i < self.rank (Batteries.UnionFind.parentD (self.findAux x).s i)

def Batteries.UnionFind.find (self : Batteries.UnionFind) (x : Fin self.size) : (s : Batteries.UnionFind) × { _root : Fin s.size // s.size = self.size }

Find root of a union-find node, updating the structure using path compression.

Equations

• self.find x = let r := self.findAux x; ⟨{ arr := r.s, parentD_lt := ⋯, rankD_lt := ⋯ }, ⟨⟨↑r.root, ⋯⟩, ⋯⟩⟩

def Batteries.UnionFind.findN {n : Nat} (self : Batteries.UnionFind) (x : Fin n) (h : n = self.size) : Batteries.UnionFind × Fin n

Find root of a union-find node, updating the structure using path compression.

Equations

• self.findN x h = match n, h, x with | .(self.size), ⋯, x => match self.find x with | ⟨s, ⟨r, h⟩⟩ => (s, Fin.cast h r)

def Batteries.UnionFind.find! (self : Batteries.UnionFind) (x : Nat) : Batteries.UnionFind × Nat

Find root of a union-find node, updating the structure using path compression. Panics if index is out of bounds.

Equations

• self.find! x = if h : x < self.size then match self.find ⟨x, h⟩ with | ⟨s, ⟨r, property⟩⟩ => (s, ↑r) else Batteries.UnionFind.panicWith (self, x) "index out of bounds"

def Batteries.UnionFind.findD (self : Batteries.UnionFind) (x : Nat) : Batteries.UnionFind × Nat

Find root of a union-find node, updating the structure using path compression. Returns inputs unchanged when index is out of bounds.

Equations

• self.findD x = if h : x < self.size then match self.find ⟨x, h⟩ with | ⟨s, ⟨r, property⟩⟩ => (s, ↑r) else (self, x)

@[simp]

theorem Batteries.UnionFind.find_size (self : Batteries.UnionFind) (x : Fin self.size) : (self.find x).fst.size = self.size

@[simp]

theorem Batteries.UnionFind.find_root_2 (self : Batteries.UnionFind) (x : Fin self.size) : ↑(self.find x).snd.val = self.rootD ↑x

@[simp]

theorem Batteries.UnionFind.find_parent_1 (self : Batteries.UnionFind) (x : Fin self.size) : (self.find x).fst.parent ↑x = self.rootD ↑x

theorem Batteries.UnionFind.find_parent_or (self : Batteries.UnionFind) (x : Fin self.size) (i : Nat) : (self.find x).fst.parent i = self.rootD i ∧ self.rootD i = self.rootD ↑x ∨ (self.find x).fst.parent i = self.parent i

@[simp]

theorem Batteries.UnionFind.find_root_1 (self : Batteries.UnionFind) (x : Fin self.size) (i : Nat) : (self.find x).fst.rootD i = self.rootD i

def Batteries.UnionFind.linkAux (self : Array Batteries.UFNode) (x : Fin self.size) (y : Fin self.size) : Array Batteries.UFNode

Link two union-find nodes

Equations

• One or more equations did not get rendered due to their size.

theorem Batteries.UnionFind.setParentBump_rankD_lt {arr' : Array Batteries.UFNode} {arr : Array Batteries.UFNode} {x : Fin arr.size} {y : Fin arr.size} (hroot : (arr.get x).rank < (arr.get y).rank ∨ (arr.get y).parent = ↑y) (H : (arr.get x).rank ≤ (arr.get y).rank) {i : Nat} (rankD_lt : Batteries.UnionFind.parentD arr i ≠ i → Batteries.UnionFind.rankD arr i < Batteries.UnionFind.rankD arr (Batteries.UnionFind.parentD arr i)) (hP : Batteries.UnionFind.parentD arr' i = if ↑x = i then ↑y else Batteries.UnionFind.parentD arr i) (hR : ∀ {i : Nat}, Batteries.UnionFind.rankD arr' i = if ↑y = i ∧ (arr.get x).rank = (arr.get y).rank then (arr.get y).rank + 1 else Batteries.UnionFind.rankD arr i) : ¬Batteries.UnionFind.parentD arr' i = i → Batteries.UnionFind.rankD arr' i < Batteries.UnionFind.rankD arr' (Batteries.UnionFind.parentD arr' i)

theorem Batteries.UnionFind.setParent_rankD_lt {arr : Array Batteries.UFNode} {x : Fin arr.size} {y : Fin arr.size} (h : (arr.get x).rank < (arr.get y).rank) {i : Nat} (rankD_lt : Batteries.UnionFind.parentD arr i ≠ i → Batteries.UnionFind.rankD arr i < Batteries.UnionFind.rankD arr (Batteries.UnionFind.parentD arr i)) : let arr' := arr.set x { parent := ↑y, rank := (arr.get x).rank }; Batteries.UnionFind.parentD arr' i ≠ i → Batteries.UnionFind.rankD arr' i < Batteries.UnionFind.rankD arr' (Batteries.UnionFind.parentD arr' i)

@[simp]

theorem Batteries.UnionFind.linkAux_size {self : Array Batteries.UFNode} {x : Fin self.size} {y : Fin self.size} : (Batteries.UnionFind.linkAux self x y).size = self.size

def Batteries.UnionFind.link (self : Batteries.UnionFind) (x : Fin self.size) (y : Fin self.size) (yroot : self.parent ↑y = ↑y) : Batteries.UnionFind

Link a union-find node to a root node.

Equations

• self.link x y yroot = { arr := Batteries.UnionFind.linkAux self.arr x y, parentD_lt := ⋯, rankD_lt := ⋯ }

def Batteries.UnionFind.linkN {n : Nat} (self : Batteries.UnionFind) (x : Fin n) (y : Fin n) (yroot : self.parent ↑y = ↑y) (h : n = self.size) : Batteries.UnionFind

Link a union-find node to a root node.

Equations

• self.linkN x y yroot h = match n, h, x, y, yroot with | .(self.size), ⋯, x, y, yroot => self.link x y yroot

def Batteries.UnionFind.link! (self : Batteries.UnionFind) (x : Nat) (y : Nat) (yroot : self.parent y = y) : Batteries.UnionFind

Link a union-find node to a root node. Panics if either index is out of bounds.

Equations

• self.link! x y yroot = if h : x < self.size ∧ y < self.size then self.link ⟨x, ⋯⟩ ⟨y, ⋯⟩ yroot else Batteries.UnionFind.panicWith self "index out of bounds"

def Batteries.UnionFind.union (self : Batteries.UnionFind) (x : Fin self.size) (y : Fin self.size) : Batteries.UnionFind

Link two union-find nodes, uniting their respective classes.

Equations

• self.union x y = match self.find x with | ⟨self₁, ⟨rx, ex⟩⟩ => let_fun hy := ⋯; match eq : self₁.find ⟨↑y, hy⟩ with | ⟨self₂, ⟨ry, ey⟩⟩ => self₂.link ⟨↑rx, ⋯⟩ ry ⋯

def Batteries.UnionFind.unionN {n : Nat} (self : Batteries.UnionFind) (x : Fin n) (y : Fin n) (h : n = self.size) : Batteries.UnionFind

Link two union-find nodes, uniting their respective classes.

Equations

• self.unionN x y h = match n, h, x, y with | .(self.size), ⋯, x, y => self.union x y

def Batteries.UnionFind.union! (self : Batteries.UnionFind) (x : Nat) (y : Nat) : Batteries.UnionFind

Link two union-find nodes, uniting their respective classes. Panics if either index is out of bounds.

Equations

• self.union! x y = if h : x < self.size ∧ y < self.size then self.union ⟨x, ⋯⟩ ⟨y, ⋯⟩ else Batteries.UnionFind.panicWith self "index out of bounds"

def Batteries.UnionFind.checkEquiv (self : Batteries.UnionFind) (x : Fin self.size) (y : Fin self.size) : Batteries.UnionFind × Bool

Check whether two union-find nodes are equivalent, updating structure using path compression.

Equations

• self.checkEquiv x y = match self.find x with | ⟨s, ⟨⟨r₁, isLt⟩, h⟩⟩ => match s.find (⋯ ▸ y) with | ⟨s_1, ⟨⟨r₂, isLt⟩, property⟩⟩ => (s_1, r₁ == r₂)

def Batteries.UnionFind.checkEquivN {n : Nat} (self : Batteries.UnionFind) (x : Fin n) (y : Fin n) (h : n = self.size) : Batteries.UnionFind × Bool

Check whether two union-find nodes are equivalent, updating structure using path compression.

Equations

• self.checkEquivN x y h = match n, h, x, y with | .(self.size), ⋯, x, y => self.checkEquiv x y

def Batteries.UnionFind.checkEquiv! (self : Batteries.UnionFind) (x : Nat) (y : Nat) : Batteries.UnionFind × Bool

Check whether two union-find nodes are equivalent, updating structure using path compression. Panics if either index is out of bounds.

Equations

• self.checkEquiv! x y = if h : x < self.size ∧ y < self.size then self.checkEquiv ⟨x, ⋯⟩ ⟨y, ⋯⟩ else Batteries.UnionFind.panicWith (self, false) "index out of bounds"

def Batteries.UnionFind.checkEquivD (self : Batteries.UnionFind) (x : Nat) (y : Nat) : Batteries.UnionFind × Bool

Check whether two union-find nodes are equivalent with path compression, returns x == y if either index is out of bounds

Equations

• self.checkEquivD x y = match self.findD x with | (s, x) => match s.findD y with | (s, y) => (s, x == y)

def Batteries.UnionFind.Equiv (self : Batteries.UnionFind) (a : Nat) (b : Nat) : Prop

Equivalence relation from a UnionFind structure

Equations

• self.Equiv a b = (self.rootD a = self.rootD b)