///|
/// Disjoint sets with union by size and path compression.
pub struct Dsu {
priv parent_or_size : Array[Int]
} derive(Debug)
///|
/// Creates n singleton sets. O(n).
pub fn Dsu::new(n : Int) -> Dsu {
guard n >= 0 && n <= 100000000 else { panic() }
{ parent_or_size: Array::make(n, -1), }
}
///|
/// Returns a representative. Amortized O(alpha(n)).
pub fn Dsu::leader(self : Dsu, a : Int) -> Int {
guard 0 <= a && a < self.parent_or_size.length() else { panic() }
let mut root = a
while self.parent_or_size[root] >= 0 {
root = self.parent_or_size[root]
}
let mut v = a
while v != root {
let next = self.parent_or_size[v]
self.parent_or_size[v] = root
v = next
}
root
}
///|
/// Unites two sets and returns their representative.
pub fn Dsu::merge(self : Dsu, a : Int, b : Int) -> Int {
let mut x = self.leader(a)
let mut y = self.leader(b)
if x == y {
return x
}
if -self.parent_or_size[x] < -self.parent_or_size[y] {
let tmp = x
x = y
y = tmp
}
self.parent_or_size[x] += self.parent_or_size[y]
self.parent_or_size[y] = x
x
}
///|
pub fn Dsu::same(self : Dsu, a : Int, b : Int) -> Bool {
self.leader(a) == self.leader(b)
}
///|
pub fn Dsu::size(self : Dsu, a : Int) -> Int {
-self.parent_or_size[self.leader(a)]
}
///|
/// Returns all nonempty groups, each in ascending vertex order. O(n).
pub fn Dsu::groups(self : Dsu) -> Array[Array[Int]] {
let buckets : Array[Array[Int]] = Array::makei(self.parent_or_size.length(), _ => {
[]
})
for i = 0; i < buckets.length(); i = i + 1 {
buckets[self.leader(i)].push(i)
}
buckets.filter(g => !g.is_empty())
}