use std::collections::BTreeMap;
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);
}
}
}