#[derive(Debug)]
struct UnionFindEntry {
parent: usize,
size: usize,
}
#[derive(Debug)]
pub struct UnionFind {
entries: Vec<UnionFindEntry>,
}
impl UnionFind {
pub fn new() -> Self {
Self {
entries: Vec::new(),
}
}
pub fn make_set(&mut self) -> usize {
let i = self.entries.len();
self.entries.push(UnionFindEntry { parent: i, size: 1 });
i
}
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;
}
}
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)
}
}