1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
use std::collections::HashMap;
use std::hash::Hash;

/// Disjoint Set Definition
pub struct DisjointSet<T>
where
    T: Clone + Eq + Hash,
{
    sets: HashMap<T, Vec<T>>,
    representatives: HashMap<T, T>,
}

impl<T> DisjointSet<T>
where
    T: Clone + Eq + Hash, // Got error coz hashmap value needs these traits
{
    /// Associate Function
    /// Creates a new Disjoint Set
    ///
    /// # Example
    ///
    /// ```
    /// let mut set = DisjointSet::new(); // creates a new Disjoint set
    /// ```
    pub fn new() -> Self {
        Self {
            sets: HashMap::new(),
            representatives: HashMap::new(),
        }
    }

    /// Insert a value to the set
    ///
    /// # Example
    ///
    /// ```
    /// // Add all the vertices of graph g to the disjoint set
    /// for (node, _) in &g.vertices {
    ///     set.set_insert(node.clone());
    /// }
    /// ```
    pub fn set_insert(&mut self, val: T) {
        self.sets.insert(val.clone(), vec![val.clone()]);
        self.representatives.insert(val.clone(), val.clone());
    }

    /// Find parent of the value
    ///
    /// # Example
    ///
    /// ```
    /// set.find(&vertex)
    /// ```
    ///
    pub fn find(&self, val: &T) -> T {
        self.representatives.get(val).unwrap().clone()
    }

    /// Union function for two nodes (vertices)
    ///
    /// # Example
    ///
    /// ```
    /// set.union(&vertex1, &vertex2)
    /// ```
    ///
    pub fn union(&mut self, a: &T, b: &T) {
        let repa = self.representatives.get(a).unwrap().clone();
        let repb = self.representatives.get(b).unwrap().clone();
        let setb = self.sets.remove(&repb).unwrap(); // get all from set of the second value

        for i in setb.iter() {
            self.representatives.remove(i); // remove them from their group
            self.representatives.insert(i.clone(), repa.clone()); // and add them to the first group
        }

        let seta = self.sets.get_mut(&repa).unwrap();

        // Now all elements from the second set will be added to first and thus union is complete
        for i in &setb {
            seta.push(i.clone());
        }
    }
}