53 lines
1.6 KiB
Text
53 lines
1.6 KiB
Text
/-
|
||
Copyright (c) 2018 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Author: Leonardo de Moura
|
||
-/
|
||
prelude
|
||
import init.data.hashmap.basic
|
||
|
||
/- Disjoint set datastructure using union-find algorithm.
|
||
We use hashmaps to implement. Thus, we should be disjoint sets
|
||
linearly for optimal performace. -/
|
||
|
||
namespace Lean
|
||
universes u
|
||
structure DisjointSet.Node (α : Type u) :=
|
||
(find : α)
|
||
(rank : Nat := 0)
|
||
|
||
structure DisjointSet (α : Type u) [HasBeq α] [Hashable α] : Type u :=
|
||
(map : HashMap α (DisjointSet.Node α))
|
||
|
||
def mkDisjointSet (α : Type u) [HasBeq α] [Hashable α] : DisjointSet α :=
|
||
⟨mkHashMap⟩
|
||
|
||
variables {α : Type u}
|
||
|
||
namespace DisjointSet
|
||
variables [DecidableEq α] [Hashable α]
|
||
|
||
private def findAux : Nat → α → HashMap α (Node α) → Node α
|
||
| 0 a m := { find := a, rank := 0 }
|
||
| (n+1) a m :=
|
||
match m.find a with
|
||
| some r := if r.find = a then r else findAux n r.find m
|
||
| none := { find := a, rank := 0 }
|
||
|
||
def find : DisjointSet α → α → α
|
||
| ⟨m⟩ a := (findAux m.size a m).find
|
||
|
||
def rank : DisjointSet α → α → Nat
|
||
| ⟨m⟩ a := (findAux m.size a m).rank
|
||
|
||
def merge : DisjointSet α → α → α → DisjointSet α
|
||
| ⟨m⟩ a b :=
|
||
let ra := findAux m.size a m in
|
||
let rb := findAux m.size b m in
|
||
if ra.find = rb.find then ⟨m⟩
|
||
else if ra.rank < rb.rank then ⟨m.insert a { find := b }⟩
|
||
else if ra.rank = rb.rank then ⟨(m.insert a { find := b }).insert b { find := b, rank := rb.rank + 1 }⟩
|
||
else ⟨m.insert b { find := a }⟩
|
||
|
||
end DisjointSet
|
||
end Lean
|