disjoint_set/
lib.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::rc::Rc;
4
5#[derive(Clone)]
6enum SetItem<ValueType> {
7    Root(ValueType, u32),
8    Link(Box<SetItem<ValueType>>)
9}
10
11
12/// Tarjan's Union-Find Data structure
13pub struct DisjointSet<T: Clone>
14{
15    set_size: usize, 
16/// The structure saves the parent information of each subset in continuous 
17/// memory(a vec) for better performance.  
18    parent: Vec<usize>,
19
20/// Each T entry is mapped onto a usize tag.
21    map: HashMap<T, usize>
22}
23
24impl<T> DisjointSet<T> where T: Clone+Hash+Eq
25{
26    pub fn new() -> Self{
27        DisjointSet{
28            set_size: 0,
29            parent: Vec::new(),
30            map: HashMap::new()
31        }
32    }
33
34    pub fn make_set(&mut self, x: T){
35        let mut len = &mut self.set_size;
36        match self.map.get(&x) {
37            Some(p) => return,
38            None => {}
39        }
40
41        self.map.insert(x, *len);
42        self.parent.push(*len);
43
44        *len += 1;
45    }
46
47    
48    pub fn find(&mut self, x: T) -> Option<usize> 
49    {
50    //! Returns Some(num), num is the tag of subset in which x is.
51    //! If x is not in the data structure, it returns None.    
52
53        let mut pos: usize;
54        match self.map.get(&x) {
55            Some(p) => {pos = *p;},
56            None => return None
57        }
58
59        let ret = DisjointSet::<T>::find_internal(&mut self.parent, pos);
60        Some(ret)            
61    }
62
63    fn find_internal(p: &mut Vec<usize>, n: usize) -> usize{
64        if p[n] != n{
65            let parent = p[n];
66            p[n] = DisjointSet::<T>::find_internal(p, parent);
67            p[n]
68        }
69        else {
70            n
71        }
72    }
73     
74   
75    pub fn union(&mut self, x: T, y: T) -> Result<usize, ()>
76    {
77        //! Union the subsets to which x and y belong.
78        //! If it returns Ok<u32>, it is the tag for unified subset.
79        //! it returns Err(), at least one of x and y is not in the disjoint-set.
80        let x_root;
81        let y_root;
82        match self.find(x) {
83            Some(x_r) => {x_root = x_r;} ,
84            None => {return Err(());}
85        }
86
87        match self.find(y) {
88            Some(y_r) => {y_root = y_r;} ,
89            None => {return Err(());}
90        }
91
92        self.parent[x_root] = y_root;
93        Ok(y_root)
94    }
95}
96
97#[test]
98fn it_works() {
99    let mut ds = DisjointSet::<i32>::new();
100    ds.make_set(1);
101    ds.make_set(2);
102    ds.make_set(3);
103
104    assert!(ds.find(1) != ds.find(2));
105    assert!(ds.find(2) != ds.find(3));
106    ds.union(1, 2);
107    ds.union(2, 3);
108    assert!(ds.find(1) == ds.find(3));
109
110    assert!(ds.find(4) == None);
111    ds.make_set(4);
112    assert!(ds.find(4) != None);
113
114
115    ds.make_set(-1);
116    assert!(ds.find(-1) != ds.find(3));
117
118    ds.union(-1, 4);
119    ds.union(2, 4);
120
121    assert!(ds.find(-1) == ds.find(3));
122}
123
124
125
disjoint_set/lib.rs

disjoint_set/
lib.rs
