pub struct UnionFind {
parent: Vec<usize>,
rank: Vec<usize>,
size: Vec<usize>,
}
impl UnionFind {
pub fn new(size: usize) -> Self {
UnionFind {
parent: (0..size).collect(),
rank: vec![0; size],
size: vec![1; size],
}
}
pub fn find(&mut self, x: usize) -> usize {
if self.parent[x] != x {
self.parent[x] = self.find(self.parent[x]);
}
self.parent[x]
}
pub fn union(&mut self, x: usize, y: usize) -> bool {
let root_x = self.find(x);
let root_y = self.find(y);
if root_x == root_y {
return false;
}
if self.rank[root_x] > self.rank[root_y] {
self.parent[root_y] = root_x;
self.size[root_x] += self.size[root_y];
} else if self.rank[root_x] < self.rank[root_y] {
self.parent[root_x] = root_y;
self.size[root_y] += self.size[root_x];
} else {
self.parent[root_y] = root_x;
self.rank[root_x] += 1;
self.size[root_x] += self.size[root_y];
}
true
}
pub fn connected(&mut self, x: usize, y: usize) -> bool {
self.find(x) == self.find(y)
}
pub fn size(&mut self, x: usize) -> usize {
let root = self.find(x);
self.size[root]
}
pub fn count_groups(&mut self) -> usize {
let n = self.parent.len();
(0..n).filter(|&i| self.find(i) == i).count()
}
}