use std::collections::{BTreeMap, HashMap};
#[derive(Debug)]
pub struct UnionFind {
parent: BTreeMap<i32, i32>,
rank: BTreeMap<i32, u32>,
}
impl UnionFind {
pub(crate) fn new() -> Self {
UnionFind {
parent: BTreeMap::new(),
rank: BTreeMap::new(),
}
}
pub(crate) fn make_set(&mut self, x: i32) {
self.parent.entry(x).or_insert(x);
self.rank.entry(x).or_insert(0);
}
pub(crate) fn find(&mut self, mut x: i32) -> i32 {
self.make_set(x);
let mut root = x;
while self.parent[&root] != root {
root = self.parent[&root];
}
while x != root {
let next = self.parent[&x];
self.parent.insert(x, root);
x = next;
}
root
}
pub(crate) fn union(&mut self, x: i32, y: i32) {
let rx = self.find(x);
let ry = self.find(y);
if rx == ry {
return;
}
let rank_x = self.rank[&rx];
let rank_y = self.rank[&ry];
if rank_x < rank_y {
self.parent.insert(rx, ry);
} else if rank_x > rank_y {
self.parent.insert(ry, rx);
} else {
self.parent.insert(ry, rx);
self.rank.insert(rx, rank_x + 1);
}
}
pub(crate) fn contains(&self, x: i32) -> bool {
self.parent.contains_key(&x)
}
pub(crate) fn ids(&self) -> impl Iterator<Item = i32> + '_ {
self.parent.keys().copied()
}
pub(crate) fn label(&mut self, ids: impl IntoIterator<Item = i32>) -> HashMap<i32, i32> {
let mut by_root: HashMap<i32, i32> = HashMap::new();
let mut labels = HashMap::new();
for id in ids {
let root = self.find(id);
let next = by_root.len() as i32;
labels.insert(id, *by_root.entry(root).or_insert(next));
}
labels
}
}