#[derive(Debug)]
pub struct UnionFindEntry {
pub parent: usize,
pub size: usize,
}
#[derive(Debug)]
pub struct UnionFind {
entries: Vec<UnionFindEntry>,
}
impl UnionFind {
pub fn with_capacity(capacity: usize) -> Self {
let mut entries = Vec::new();
for i in 0..capacity {
entries.push(UnionFindEntry { parent: i, size: 1 });
}
Self { entries }
}
pub fn find(&self, x: usize) -> usize {
let mut root = x;
loop {
let parent = self.entries[root].parent;
if parent == root {
break;
}
root = parent;
}
root
}
fn find_mut(&mut self, mut x: usize) -> usize {
let mut root = x;
loop {
let parent = self.entries[root].parent;
if parent == root {
break;
}
root = parent;
}
loop {
let entry = &mut self.entries[x];
if entry.parent == root {
break;
}
let parent = entry.parent;
entry.parent = root;
x = parent;
}
root
}
pub fn union(&mut self, mut x: usize, mut y: usize) {
x = self.find_mut(x);
y = self.find_mut(y);
if x == y {
return;
}
let x_size = self.entries[x].size;
let y_size = self.entries[y].size;
if x_size > y_size {
self.entries[y].parent = x;
self.entries[x].size += y_size;
} else {
self.entries[x].parent = y;
self.entries[y].size += x_size;
}
}
pub fn leaders(&self) -> impl Iterator<Item = &UnionFindEntry> {
self.entries.iter().enumerate().filter_map(|(i, entry)| {
if i != entry.parent {
None
} else {
Some(entry)
}
})
}
pub fn largest(&self) -> Option<usize> {
let mut max: Option<&UnionFindEntry> = None;
for entry in self.leaders() {
match max {
None => {
max = Some(entry);
}
Some(current_entry) => {
if entry.size > current_entry.size {
max = Some(entry);
}
}
}
}
max.map(|entry| entry.parent)
}
pub fn size(&self, index: usize) -> usize {
let entry = &self.entries[index];
debug_assert!(entry.parent == index);
entry.size
}
pub fn sizes(&self) -> impl Iterator<Item = usize> + '_ {
self.leaders().map(|entry| entry.size)
}
}