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 disjoint_set.node (α : Type u) :=
|
||
(find : α)
|
||
(rank : nat := 0)
|
||
|
||
structure disjoint_set (α : Type u) [decidable_eq α] [hashable α] : Type u :=
|
||
(map : hashmap α (disjoint_set.node α))
|
||
|
||
def mk_disjoint_set (α : Type u) [decidable_eq α] [hashable α] : disjoint_set α :=
|
||
⟨mk_hashmap⟩
|
||
|
||
variables {α : Type u}
|
||
|
||
namespace disjoint_set
|
||
variables [decidable_eq α] [hashable α]
|
||
|
||
private def find_aux : 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 find_aux n r.find m
|
||
| none := { find := a, rank := 0 }
|
||
|
||
def find : disjoint_set α → α → α
|
||
| ⟨m⟩ a := (find_aux m.size a m).find
|
||
|
||
def rank : disjoint_set α → α → nat
|
||
| ⟨m⟩ a := (find_aux m.size a m).rank
|
||
|
||
def merge : disjoint_set α → α → α → disjoint_set α
|
||
| ⟨m⟩ a b :=
|
||
let ra := find_aux m.size a m in
|
||
let rb := find_aux 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 disjoint_set
|
||
end lean
|