///|
/// A simple union-find with path compression and union by size.
pub struct UnionFind {
  parents : Array[Id]
  sizes : Array[Int]
} derive(Show)

///|
pub fn UnionFind::new() -> UnionFind {
  UnionFind::{ parents: [], sizes: [] }
}

///|
pub fn UnionFind::make_set(self : UnionFind) -> Id {
  let id = self.parents.length()
  self.parents.push(id)
  self.sizes.push(1)
  id
}

///|
pub fn UnionFind::find(self : UnionFind, id : Id) -> Id {
  let parent = self.parents[id]
  if parent == id {
    id
  } else {
    let root = self.find(parent)
    self.parents[id] = root
    root
  }
}

///|
pub fn UnionFind::find_read(self : UnionFind, id : Id) -> Id {
  loop id {
    root if self.parents[root] == root => root
    node => continue self.parents[node]
  }
}

///|
pub fn UnionFind::union(self : UnionFind, a : Id, b : Id) -> Id {
  let mut root_a = self.find(a)
  let mut root_b = self.find(b)
  if root_a == root_b {
    return root_a
  }
  if self.sizes[root_a] < self.sizes[root_b] {
    let tmp = root_a
    root_a = root_b
    root_b = tmp
  }
  self.parents[root_b] = root_a
  self.sizes[root_a] = self.sizes[root_a] + self.sizes[root_b]
  root_a
}