Rust——并查集
实现
#![allow(unused)]
//! 并查集:解决节点连接/关联问题
/// 并查集 trait
trait UF {
fn is_connected(&mut self, p: usize, q: usize) -> bool;
fn union_elements(&mut self, p: usize, q: usize);
fn get_size(&self) -> usize;
}
/// 并查集结构
pub struct UnionFind {
// 存储父节点
parent: Vec<usize>,
// 高度
rank: Vec<usize>,
}
impl UnionFind {
/// 构造
pub fn new_with_size(size: usize) -> Self {
let mut res = Self {
parent: vec![0_usize; size],
rank: vec![1_usize; size],
};
// 为每个元素分组
for i in 0..res.parent.len() {
res.parent[i] = i;
}
return res;
}
/// 查询p的根
fn find(&mut self, p: usize) -> Result<usize, &'static str> {
if p >= self.parent.len() {
return Err("参数错误");
}
let mut c = p;
// 寻找根
while c != self.parent[c] {
// 压缩高度
self.parent[c] = self.parent[self.parent[c]];
c = self.parent[c];
}
return Ok(c);
}
}
impl UF for UnionFind {
/// 查看两元素是不是同一个根
fn is_connected(&mut self, p: usize, q: usize) -> bool {
let p_root = self.find(p).unwrap();
let q_root = self.find(q).unwrap();
return p_root == q_root;
}
/// 合并两个元素为一个根
fn union_elements(&mut self, p: usize, q: usize) {
let p_root = self.find(p).unwrap();
let q_root = self.find(q).unwrap();
if p_root != q_root {
if self.rank[p_root] < self.rank[q_root] {
self.parent[p_root] = self.parent[q_root];
} else if self.rank[q_root] < self.rank[p_root] {
self.parent[q_root] = self.parent[p_root];
} else {
self.parent[q_root] = self.parent[p_root];
self.rank[p_root] += 1;
}
}
}
fn get_size(&self) -> usize {
return self.parent.len();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_works() {
let mut union_find = UnionFind::new_with_size(10);
union_find.union_elements(3, 5);
union_find.union_elements(2, 1);
union_find.union_elements(5, 1);
union_find.union_elements(5, 4);
assert_eq!(union_find.is_connected(4, 1), true);
}
}