71 lines
2.1 KiB
Text
71 lines
2.1 KiB
Text
namespace Batteries
|
|
|
|
/-- Union-find node type -/
|
|
structure UFNode where
|
|
/-- Parent of node -/
|
|
parent : Nat
|
|
|
|
namespace UnionFind
|
|
|
|
/-- Parent of a union-find node, defaults to self when the node is a root -/
|
|
def parentD (arr : Array UFNode) (i : Nat) : Nat :=
|
|
if h : i < arr.size then arr[i].parent else i
|
|
|
|
/-- Rank of a union-find node, defaults to 0 when the node is a root -/
|
|
def rankD (arr : Array UFNode) (i : Nat) : Nat := 0
|
|
|
|
theorem parentD_of_not_lt : ¬i < arr.size → parentD arr i = i := sorry
|
|
|
|
theorem parentD_set {arr : Array UFNode} {x h v i} :
|
|
parentD (arr.set x v h) i = if x = i then v.parent else parentD arr i := by
|
|
rw [parentD]
|
|
sorry
|
|
|
|
end UnionFind
|
|
|
|
open UnionFind
|
|
|
|
structure UnionFind where
|
|
arr : Array UFNode
|
|
|
|
namespace UnionFind
|
|
|
|
/-- Size of union-find structure. -/
|
|
@[inline] abbrev size (self : UnionFind) := self.arr.size
|
|
|
|
/-- Parent of union-find node -/
|
|
abbrev parent (self : UnionFind) (i : Nat) : Nat := parentD self.arr i
|
|
|
|
theorem parent_lt (self : UnionFind) (i : Nat) : self.parent i < self.size ↔ i < self.size :=
|
|
sorry
|
|
|
|
/-- Rank of union-find node -/
|
|
abbrev rank (self : UnionFind) (i : Nat) : Nat := rankD self.arr i
|
|
|
|
/-- Maximum rank of nodes in a union-find structure -/
|
|
noncomputable def rankMax (self : UnionFind) := 0
|
|
|
|
/-- Root of a union-find node. -/
|
|
def root (self : UnionFind) (x : Fin self.size) : Fin self.size :=
|
|
let y := self.arr[x.1].parent
|
|
if h : y = x then
|
|
x
|
|
else
|
|
have : self.rankMax - self.rank self.arr[x.1].parent < self.rankMax - self.rank x :=
|
|
sorry
|
|
self.root ⟨y, sorry⟩
|
|
termination_by self.rankMax - self.rank x
|
|
|
|
/-- Root of a union-find node. Returns input if index is out of bounds. -/
|
|
def rootD (self : UnionFind) (x : Nat) : Nat :=
|
|
if h : x < self.size then self.root ⟨x, h⟩ else x
|
|
|
|
theorem rootD_parent (self : UnionFind) (x : Nat) : self.rootD (self.parent x) = self.rootD x := by
|
|
simp only [rootD, Array.length_toList, parent_lt]
|
|
split
|
|
· simp only [parentD, ↓reduceDIte, *]
|
|
conv => rhs; rw [root]
|
|
split
|
|
· rw [root, dif_pos] <;> simp_all
|
|
· simp
|
|
· simp only [not_false_eq_true, parentD_of_not_lt, *]
|