pub struct Dsu {
parent: Vec<usize>,
size: Vec<usize>,
components: usize,
}
impl Dsu {
#[must_use]
pub fn new(n: usize) -> Self {
Self {
parent: (0..n).collect(),
size: vec![1; n],
components: n,
}
}
#[inline]
#[must_use]
pub fn parent(&self) -> &[usize] {
&self.parent
}
#[inline]
#[must_use]
pub fn size(&self) -> &[usize] {
&self.size
}
#[inline]
#[must_use]
pub const fn components(&self) -> usize {
self.components
}
#[inline]
pub fn find(&mut self, mut x: usize) -> usize {
let mut root = x;
while self.parent[root] != root {
root = self.parent[root];
}
while self.parent[x] != x {
let next = self.parent[x];
self.parent[x] = root;
x = next;
}
root
}
#[inline]
pub fn union(&mut self, x: usize, y: usize) -> bool {
let rx = self.find(x);
let ry = self.find(y);
if rx == ry {
return false;
}
if self.size[rx] < self.size[ry] {
self.parent[rx] = ry;
self.size[ry] += self.size[rx];
} else {
self.parent[ry] = rx;
self.size[rx] += self.size[ry];
}
self.components -= 1;
true
}
}
#[cfg(test)]
mod tests {
use super::Dsu;
#[test]
fn new() {
let input = 8;
let expected = (vec![0, 1, 2, 3, 4, 5, 6, 7], vec![1; 8], 8);
let dsu = Dsu::new(input);
let output = (dsu.parent, dsu.size, dsu.components);
assert_eq!(expected, output, "\n input: {input:?}");
}
#[test]
fn find() {
let input = [0, 3, 1, 2];
let expected = [0, 3, 0, 0];
let mut dsu = Dsu {
parent: vec![0, 0, 1, 3],
size: vec![2, 1, 1, 1],
components: 2,
};
let output = input.map(|n| dsu.find(n));
assert_eq!(expected, output, "\n input: {input:?}");
}
#[test]
fn find_compression() {
let input = 2;
let expected = 0;
let mut dsu = Dsu {
parent: vec![0, 0, 1, 3],
size: vec![2, 1, 1, 1],
components: 2,
};
let _ = dsu.find(input);
let output = dsu.parent()[2];
assert_eq!(expected, output, "\n input: {input:?}");
}
#[test]
fn union_size() {
let input = [((0, 1), 0), ((2, 0), 0), ((3, 0), 0), ((3, 1), 0)];
let expected = [(true, 2), (true, 3), (true, 4), (false, 4)];
let mut dsu = Dsu::new(4);
let output = input.map(|((x, y), r)| {
let union = dsu.union(x, y);
let root = dsu.find(r);
(union, dsu.size()[root])
});
assert_eq!(expected, output, "\n input: {input:?}");
}
#[test]
fn union_components() {
let input = [(0, 1), (2, 3), (1, 3), (0, 2)];
let expected = [(true, 3), (true, 2), (true, 1), (false, 1)];
let mut dsu = Dsu::new(4);
let output = input.map(|(x, y)| (dsu.union(x, y), dsu.components()));
assert_eq!(expected, output, "\n input: {input:?}");
}
}