///|
priv struct UnionFind {
parents : Array[Int]
ranks : Array[Int]
}
///|
fn UnionFind::validate_union_find_index(
self : UnionFind,
index : Int,
) -> Result[Unit, DedupError] {
if index < 0 || index >= self.parents.length() {
Err(UnionFindIndexOutOfBounds(index, self.parents.length()))
} else {
Ok(())
}
}
///|
fn UnionFind::new(size : Int) -> Result[UnionFind, DedupError] {
if size < 0 {
return Err(InvalidUnionFindSize(size))
}
let parents = Array::makei(size, fn(index) { index })
Ok({ parents, ranks: Array::make(size, 0) })
}
///|
fn UnionFind::length(self : UnionFind) -> Int {
self.parents.length()
}
///|
fn UnionFind::find(self : UnionFind, index : Int) -> Result[Int, DedupError] {
match self.validate_union_find_index(index) {
Err(error) => return Err(error)
Ok(_) => ()
}
let mut root = index
while self.parents[root] != root {
root = self.parents[root]
}
let mut current = index
while self.parents[current] != current {
let next = self.parents[current]
self.parents[current] = root
current = next
}
Ok(root)
}
///|
fn UnionFind::union(
self : UnionFind,
left : Int,
right : Int,
) -> Result[Int, DedupError] {
let left_root = match self.find(left) {
Err(error) => return Err(error)
Ok(value) => value
}
let right_root = match self.find(right) {
Err(error) => return Err(error)
Ok(value) => value
}
if left_root == right_root {
return Ok(left_root)
}
let left_rank = self.ranks[left_root]
let right_rank = self.ranks[right_root]
if left_rank > right_rank {
self.parents[right_root] = left_root
return Ok(left_root)
}
if right_rank > left_rank {
self.parents[left_root] = right_root
return Ok(right_root)
}
let preferred = if left_root < right_root { left_root } else { right_root }
let other = if preferred == left_root { right_root } else { left_root }
self.parents[other] = preferred
self.ranks[preferred] = self.ranks[preferred] + 1
Ok(preferred)
}