luka/
dsu.rs

1use std::fmt;
2
3#[allow(clippy::upper_case_acronyms)]
4pub struct DSU {
5    parents: Vec::<usize>,
6    ranks: Vec::<usize>
7}
8
9#[allow(clippy::upper_case_acronyms)]
10#[derive(Debug, Clone)]
11pub struct DSUError {
12    msg: String
13}
14
15impl DSUError {
16    fn new(msg: &str) -> Self {
17        DSUError{
18            msg: msg.to_string()
19        }
20    }
21}
22
23impl fmt::Display for DSUError {
24    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
25        write!(f, "{}", self.msg)
26    }
27}
28
29impl DSU {
30    pub fn new(n: usize) -> Self {
31        DSU {
32            parents: vec![0; n],
33            ranks: vec![1; n]
34        }
35    }
36    pub fn make_set(&mut self, value: usize) -> Result<(), DSUError> {
37        if value == 0 {
38            return Err(DSUError::new("Cannot create item with value 0"));
39        }
40        if value < self.parents.len() {
41            self.parents[value] = value;
42            return Ok(());
43        }
44        Err(DSUError::new("Cannot create item exceeds set size"))
45    }
46
47    pub fn find_set(&mut self, value: usize) -> Option<usize> {
48        if value >= self.parents.len() || value == 0 {
49            return None;
50        }
51        if value == self.parents[value] {
52            return Some(value);
53        }
54        match self.find_set(self.parents[value]) {
55            Some(next) => {
56                self.parents[value] = next;
57                match next {
58                    0 => None,
59                    _ => Some(next)
60                }
61            }
62            None => None
63        }
64    }
65
66    pub fn union_sets(&mut self, first: usize, second: usize) -> Result<(), DSUError>{
67        let first = self.find_set(first);
68        let second = self.find_set(second);
69        if first.is_none() || second.is_none() {
70            return Err(DSUError::new("Some elements are missing from the set. Unification is not possible"));
71        }
72        let first = first.unwrap();
73        let second = second.unwrap();
74        if first != second {
75            if self.ranks[first] < self.ranks[second] {
76                self.ranks.swap(first, second);
77            }
78            self.parents[second] = first;
79            if self.ranks[first] == self.ranks[second] {
80                self.ranks[first] += 1;
81            }
82        }
83        Ok(())
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_dsu() {
93        let mut dsu = DSU::new(10);
94
95        dsu.make_set(1).unwrap();
96        dsu.make_set(2).unwrap();
97        dsu.make_set(3).unwrap();
98        dsu.make_set(7).unwrap();
99        dsu.make_set(8).unwrap();
100
101        dsu.union_sets(1, 2).unwrap();
102        dsu.union_sets(2, 3).unwrap();
103        dsu.union_sets(2, 7).unwrap();
104
105        assert_eq!(dsu.union_sets(0, 7).is_err(), true);
106        assert_eq!(dsu.find_set(2).unwrap(), dsu.find_set(7).unwrap());
107        assert_eq!(dsu.find_set(1).unwrap(), dsu.find_set(3).unwrap());
108        assert_ne!(dsu.find_set(1).unwrap(), dsu.find_set(8).unwrap());
109        assert_eq!(dsu.find_set(29), None);
110        assert_eq!(dsu.union_sets(9, 2).is_err(), true);
111    }
112
113    #[test]
114    #[should_panic]
115    fn test_dsu_panic() {
116        let mut dsu = DSU::new(10);
117        dsu.make_set(77).unwrap();
118        assert!(dsu.make_set(77).is_err())
119    }
120}
luka/dsu.rs

luka/
dsu.rs
