use std::fs;
use super::Graph;
pub struct UF {
id: Vec<usize>,
sz: Vec<usize>,
count: usize,
}
impl UF {
pub fn new(n: usize) -> Self {
let mut id = vec![0; n];
let sz = vec![1; n];
(0..n).for_each(|i| {
id[i] = i;
});
UF { id, sz, count: n }
}
pub fn count(&self) -> usize {
self.count
}
pub fn connected(&self, p: usize, q: usize) -> bool {
self.find(p) == self.find(q)
}
pub fn find(&self, p: usize) -> usize {
let mut p = p;
while p != self.id[p] {
p = self.id[p];
}
p
}
pub fn union(&mut self, p: usize, q: usize) {
let i = self.find(p);
let j = self.find(q);
if i != j {
if self.sz[i] < self.sz[j] {
self.id[i] = j;
self.sz[j] += self.sz[i];
} else {
self.id[j] = i;
self.sz[i] += self.sz[j];
}
self.count -= 1;
}
}
}
pub struct Connect {
marked: Vec<bool>,
id: Vec<usize>,
count: usize,
}
impl Connect {
pub fn new(g: &Graph) -> Self {
let v = g.vertex_count();
let mut marked = Vec::with_capacity(v);
let mut id = Vec::with_capacity(v);
for _ in 0..v {
marked.push(false);
id.push(usize::MAX);
}
let mut connect = Connect {
marked,
id,
count: 0,
};
for s in 0..v {
if !connect.marked[s] {
connect.dfs(g, s);
connect.count += 1;
}
}
connect
}
fn dfs(&mut self, g: &Graph, v: usize) {
self.marked[v] = true;
self.id[v] = self.count;
for w in g.adj(v).clone() {
if !self.marked[w] {
self.dfs(g, w);
}
}
}
pub fn connected(&self, v: usize, w: usize) -> bool {
self.id[v] == self.id[w]
}
pub fn id(&self, v: usize) -> usize {
self.id[v]
}
pub fn count(&self) -> usize {
self.count
}
pub fn print_dot(&self, path: &str) {
let mut graph = r#"graph{
label="graph"
bgcolor=grey
shape=circle
"#
.to_string();
for t in 0..self.count {
graph.push_str(&format!(
"trees_{}[color=red;style=filled;label=\"{}\"]\n",
t, t
));
}
for (v, w) in self.id.iter().enumerate() {
graph.push_str(&format!("{} -- trees_{}\n", v, w));
}
graph.push('}');
let _ = fs::write(path, graph);
}
}
#[cfg(test)]
mod test {
use crate::{
add_edge,
graph::{connect::Connect, Graph},
};
use super::UF;
#[test]
fn test() {
let mut g = Graph::new(13);
add_edge!(g, 0, 6, 5, 2, 1);
g.add_edge(1, 0);
g.add_edge(2, 0);
add_edge!(g, 3, 5, 4);
add_edge!(g, 4, 5, 6, 3);
add_edge!(g, 5, 3, 4, 0);
add_edge!(g, 6, 0, 4);
add_edge!(g, 7, 8);
add_edge!(g, 8, 7);
add_edge!(g, 9, 11, 10, 12);
add_edge!(g, 10, 9);
add_edge!(g, 11, 9, 12);
add_edge!(g, 12, 11, 9);
let connect = Connect::new(&g);
assert_eq!(connect.count(), 3);
assert_eq!(connect.connected(8, 7), true);
assert_eq!(connect.connected(9, 10), true);
assert_eq!(connect.connected(0, 6), true);
assert_eq!(connect.connected(0, 8), false);
assert_eq!(connect.connected(10, 12), true);
assert_eq!(connect.id(6), 0);
assert_eq!(connect.id(8), 1);
assert_eq!(connect.id(12), 2);
}
#[test]
fn test_union_find() {
let mut uf = UF::new(8);
assert_eq!(uf.count(), 8);
uf.union(0, 1);
assert_eq!(uf.count(), 7);
uf.union(2, 3);
assert_eq!(uf.count(), 6);
uf.union(4, 5);
assert_eq!(uf.count(), 5);
uf.union(6, 7);
assert_eq!(uf.count(), 4);
assert!(uf.connected(0, 1));
assert!(!uf.connected(2, 1));
uf.union(0, 2);
uf.union(4, 6);
uf.union(0, 4);
assert_eq!(uf.count(), 1);
}
}