1use std::collections::BTreeMap;
2
3pub struct DSURef<'a, T> where T: Eq + Ord {
47 parent: BTreeMap<&'a T, &'a T>,
48 ranks: BTreeMap<&'a T, usize>
49}
50
51impl <'a, T> DSURef<'a, T> where T: Eq + Ord {
52 pub fn new() -> Self {
53 DSURef {
54 parent: BTreeMap::new(),
55 ranks: BTreeMap::new()
56 }
57 }
58 pub fn make_set(&mut self, value: &'a T) {
59 if !self.parent.contains_key(&value) {
60 self.parent.insert(value, value);
61 self.ranks.insert(value, 1);
62 }
63 }
64
65 pub fn find_set(&mut self, value: &'a T) -> Option<&'a T> {
66 if !self.parent.contains_key(value) {
67 return None;
68 }
69 let reference = *self.parent.get(value).unwrap();
70 if *value == *reference {
71 return Some(value);
72 }
73 let next = self.find_set(reference).unwrap();
74 *self.parent.get_mut(value).unwrap() = next;
75 return Some(next);
76 }
77
78 pub fn union_sets(&mut self, first: &'a T, second: &'a T) {
79 let first = self.find_set(first);
80 let second = self.find_set(second);
81 if first.is_some() && second.is_some() {
82 if *first.unwrap() != *second.unwrap() {
83 let first_rank = *self.ranks.get(first.as_ref().unwrap()).unwrap();
84 let second_rank = *self.ranks.get(second.as_ref().unwrap()).unwrap();
85 if second_rank >= first_rank {
86 let key = second.unwrap();
87 *self.parent.get_mut(&key).unwrap() = first.unwrap();
88 *self.ranks.get_mut(&key).unwrap() = first_rank + second_rank;
89 } else {
90 let key = first.unwrap();
91 *self.parent.get_mut(&key).unwrap() = second.unwrap();
92 *self.ranks.get_mut(&key).unwrap() = first_rank + second_rank;
93 }
94 }
95 }
96 }
97}
98
99pub struct DSU<T> where T: Eq + Ord + Clone {
100 parent: BTreeMap<T, T>,
101 ranks: BTreeMap<T, usize>
102}
103
104impl <T> DSU<T> where T: Eq + Ord + Clone {
105 pub fn new() -> Self {
106 DSU {
107 parent: BTreeMap::new(),
108 ranks: BTreeMap::new()
109 }
110 }
111 pub fn make_set(&mut self, value: T) {
112 if !self.parent.contains_key(&value) {
113 self.parent.insert(value.clone(), value.clone());
114 self.ranks.insert(value.clone(), 1);
115 }
116 }
117
118 pub fn find_set(&mut self, value: T) -> Option<T> {
119 if !self.parent.contains_key(&value) {
120 return None;
121 }
122 let reference = self.parent.get(&value).unwrap().clone();
123 if value == reference {
124 return Some(value);
125 }
126 let next = self.find_set(reference).unwrap();
127 *self.parent.get_mut(&value).unwrap() = next.clone();
128 return Some(next);
129 }
130
131 pub fn union_sets(&mut self, first: T, second: T) {
132 let first = self.find_set(first);
133 let second = self.find_set(second);
134 if first.is_some() && second.is_some() {
135 if first.as_ref().unwrap() != second.as_ref().unwrap() {
136 let first_rank = *self.ranks.get(&first.as_ref().unwrap()).unwrap();
137 let second_rank = *self.ranks.get(&second.as_ref().unwrap()).unwrap();
138 if second_rank >= first_rank {
139 let key = second.unwrap();
140 *self.parent.get_mut(&key).unwrap() = first.unwrap();
141 *self.ranks.get_mut(&key).unwrap() = first_rank + second_rank;
142 } else {
143 let key = first.unwrap();
144 *self.parent.get_mut(&key).unwrap() = second.unwrap();
145 *self.ranks.get_mut(&key).unwrap() = first_rank + second_rank;
146 }
147 }
148 }
149 }
150}
151
152pub struct DSUNum {
153 parent: Vec::<usize>,
154 ranks: Vec::<usize>
155}
156
157impl DSUNum {
158 pub fn new(n: usize) -> Self {
159 DSUNum {
160 parent: vec![0; n + 1],
161 ranks: vec![1; n + 1]
162 }
163 }
164 pub fn make_set(&mut self, value: usize) {
165 self.parent[value] = value;
166 }
167
168 pub fn find_set(&mut self, value: usize) -> usize {
169 if value == self.parent[value] {
170 return value;
171 }
172 let next = self.find_set(self.parent[value]);
173 self.parent[value] = next;
174 return next;
175 }
176
177 pub fn union_sets(&mut self, first: usize, second: usize) {
178 let first = self.find_set(first);
179 let second = self.find_set(second);
180 if first != second {
181 if self.ranks[first] < self.ranks[second] {
182 self.parent[second] = first;
183 self.ranks[second] += self.ranks[first];
184 } else {
185 self.parent[first] = second;
186 self.ranks[first] += self.ranks[second];
187 }
188 }
189 }
190}
191
192#[test]
193fn test_dsu_ref() {
194 let mut dsu = DSURef::new();
195 let v = (0..10).collect::<Vec<u32>>();
196 for i in &v {
197 dsu.make_set(i);
198 }
199 dsu.union_sets(&v[1], &v[2]);
200 dsu.union_sets(&v[2], &v[3]);
201 dsu.union_sets(&v[2], &v[7]);
202
203 assert_eq!(dsu.find_set(&v[2]), dsu.find_set(&v[7]));
204 assert_ne!(dsu.find_set(&v[2]), dsu.find_set(&v[8]));
205 assert_eq!(dsu.find_set(&v[1]), dsu.find_set(&v[3]));
206 assert_ne!(dsu.find_set(&v[1]), dsu.find_set(&v[9]));
207 assert_eq!(dsu.find_set(&11), None);
208}
209
210#[test]
211fn test_dsu() {
212 let mut dsu = DSU::new();
213 for i in 0..10 {
214 dsu.make_set(i);
215 }
216 dsu.union_sets(1, 2);
217 dsu.union_sets(2, 3);
218 dsu.union_sets(2, 7);
219
220 assert_eq!(dsu.find_set(2), dsu.find_set(7));
221 assert_ne!(dsu.find_set(2), dsu.find_set(8));
222 assert_eq!(dsu.find_set(1), dsu.find_set(3));
223 assert_ne!(dsu.find_set(1), dsu.find_set(9));
224 assert_eq!(dsu.find_set(11), None);
225}
226
227#[test]
228fn test_dsu_num() {
229 let mut dsu = DSUNum::new(10);
230 for i in 1..=10 {
231 dsu.make_set(i);
232 }
233 dsu.union_sets(1, 2);
234 dsu.union_sets(2, 3);
235 dsu.union_sets(2, 7);
236
237 assert_eq!(dsu.find_set(2), dsu.find_set(7));
238 assert_ne!(dsu.find_set(2), dsu.find_set(8));
239 assert_eq!(dsu.find_set(1), dsu.find_set(3));
240 assert_ne!(dsu.find_set(1), dsu.find_set(9));
241}