1use std::fs;
60
61use super::Graph;
62pub struct UF {
63 id: Vec<usize>,
65 sz: Vec<usize>,
67 count: usize,
69}
70impl UF {
71 pub fn new(n: usize) -> Self {
72 let mut id = vec![0; n];
73 let sz = vec![1; n];
74 (0..n).for_each(|i| {
75 id[i] = i;
76 });
77 UF { id, sz, count: n }
78 }
79 pub fn count(&self) -> usize {
80 self.count
81 }
82 pub fn connected(&self, p: usize, q: usize) -> bool {
83 self.find(p) == self.find(q)
84 }
85
86 pub fn find(&self, p: usize) -> usize {
87 let mut p = p;
88 while p != self.id[p] {
89 p = self.id[p];
90 }
91 p
92 }
93 pub fn union(&mut self, p: usize, q: usize) {
94 let i = self.find(p);
95 let j = self.find(q);
96 if i != j {
97 if self.sz[i] < self.sz[j] {
98 self.id[i] = j;
99 self.sz[j] += self.sz[i];
100 } else {
101 self.id[j] = i;
102 self.sz[i] += self.sz[j];
103 }
104 self.count -= 1;
105 }
106 }
107}
108pub struct Connect {
109 marked: Vec<bool>,
110 id: Vec<usize>,
111 count: usize,
112}
113impl Connect {
114 pub fn new(g: &Graph) -> Self {
115 let v = g.vertex_count();
116 let mut marked = Vec::with_capacity(v);
117 let mut id = Vec::with_capacity(v);
118 for _ in 0..v {
119 marked.push(false);
120 id.push(usize::MAX);
121 }
122 let mut connect = Connect {
123 marked,
124 id,
125 count: 0,
126 };
127 for s in 0..v {
128 if !connect.marked[s] {
129 connect.dfs(g, s);
130 connect.count += 1;
131 }
132 }
133 connect
134 }
135
136 fn dfs(&mut self, g: &Graph, v: usize) {
137 self.marked[v] = true;
138 self.id[v] = self.count;
139 for w in g.adj(v).clone() {
140 if !self.marked[w] {
141 self.dfs(g, w);
142 }
143 }
144 }
145 pub fn connected(&self, v: usize, w: usize) -> bool {
146 self.id[v] == self.id[w]
147 }
148 pub fn id(&self, v: usize) -> usize {
149 self.id[v]
150 }
151 pub fn count(&self) -> usize {
152 self.count
153 }
154 pub fn print_dot(&self, path: &str) {
156 let mut graph = r#"graph{
157 label="graph"
158 bgcolor=grey
159 shape=circle
160 "#
161 .to_string();
162 for t in 0..self.count {
163 graph.push_str(&format!(
164 "trees_{}[color=red;style=filled;label=\"{}\"]\n",
165 t, t
166 ));
167 }
168 for (v, w) in self.id.iter().enumerate() {
169 graph.push_str(&format!("{} -- trees_{}\n", v, w));
170 }
171 graph.push('}');
172 let _ = fs::write(path, graph);
173 }
174}
175
176#[cfg(test)]
177mod test {
178 use crate::{
179 add_edge,
180 graph::{connect::Connect, Graph},
181 };
182
183 use super::UF;
184
185 #[test]
186 fn test() {
187 let mut g = Graph::new(13);
188 add_edge!(g, 0, 6, 5, 2, 1);
189 g.add_edge(1, 0);
190 g.add_edge(2, 0);
191 add_edge!(g, 3, 5, 4);
192 add_edge!(g, 4, 5, 6, 3);
193 add_edge!(g, 5, 3, 4, 0);
194 add_edge!(g, 6, 0, 4);
195 add_edge!(g, 7, 8);
196 add_edge!(g, 8, 7);
197 add_edge!(g, 9, 11, 10, 12);
198 add_edge!(g, 10, 9);
199 add_edge!(g, 11, 9, 12);
200 add_edge!(g, 12, 11, 9);
201 let connect = Connect::new(&g);
203 assert_eq!(connect.count(), 3);
205 assert_eq!(connect.connected(8, 7), true);
206 assert_eq!(connect.connected(9, 10), true);
207 assert_eq!(connect.connected(0, 6), true);
208 assert_eq!(connect.connected(0, 8), false);
209 assert_eq!(connect.connected(10, 12), true);
210 assert_eq!(connect.id(6), 0);
211 assert_eq!(connect.id(8), 1);
212 assert_eq!(connect.id(12), 2);
213 }
214 #[test]
215 fn test_union_find() {
216 let mut uf = UF::new(8);
217 assert_eq!(uf.count(), 8);
218 uf.union(0, 1);
219 assert_eq!(uf.count(), 7);
220 uf.union(2, 3);
221 assert_eq!(uf.count(), 6);
222 uf.union(4, 5);
223 assert_eq!(uf.count(), 5);
224 uf.union(6, 7);
225 assert_eq!(uf.count(), 4);
226 assert!(uf.connected(0, 1));
227 assert!(!uf.connected(2, 1));
228 uf.union(0, 2);
229 uf.union(4, 6);
230 uf.union(0, 4);
231 assert_eq!(uf.count(), 1);
232 }
233}