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);
    }
}