1use crate::dulmage_mendelsohn::{DMPart, DulmageMendelsohnPartition};
9use crate::incidence::EqualityIncidence;
10use crate::matching::BipartiteMatching;
11
12#[derive(Debug, Clone, Default)]
26pub struct SquareComponent {
27 pub eq_rows: Vec<usize>,
30 pub cols: Vec<usize>,
32}
33
34#[derive(Debug, Clone, Default)]
36pub struct SquareComponents {
37 pub components: Vec<SquareComponent>,
40}
41
42struct UnionFind {
45 parent: Vec<usize>,
46 rank: Vec<u8>,
47}
48
49impl UnionFind {
50 fn new(n: usize) -> Self {
51 Self {
52 parent: (0..n).collect(),
53 rank: vec![0; n],
54 }
55 }
56
57 fn find(&mut self, x: usize) -> usize {
58 let mut root = x;
59 while self.parent[root] != root {
60 root = self.parent[root];
61 }
62 let mut cur = x;
64 while self.parent[cur] != root {
65 let nxt = self.parent[cur];
66 self.parent[cur] = root;
67 cur = nxt;
68 }
69 root
70 }
71
72 fn union(&mut self, x: usize, y: usize) {
73 let rx = self.find(x);
74 let ry = self.find(y);
75 if rx == ry {
76 return;
77 }
78 match self.rank[rx].cmp(&self.rank[ry]) {
79 std::cmp::Ordering::Less => self.parent[rx] = ry,
80 std::cmp::Ordering::Greater => self.parent[ry] = rx,
81 std::cmp::Ordering::Equal => {
82 self.parent[ry] = rx;
83 self.rank[rx] += 1;
84 }
85 }
86 }
87}
88
89impl SquareComponents {
90 pub fn of_square_part(
124 inc: &EqualityIncidence,
125 _m: &BipartiteMatching,
126 part: &DulmageMendelsohnPartition,
127 ) -> Self {
128 let n_rows = inc.n_eq_rows();
129 let n_vars = inc.n_vars;
130
131 let mut uf = UnionFind::new(n_rows + n_vars);
133 let col_off = n_rows;
134
135 for r in 0..n_rows {
136 if part.row_part[r] != DMPart::Square {
137 continue;
138 }
139 for &v in inc.neighbors(r) {
140 if part.col_part[v] != DMPart::Square {
141 continue;
142 }
143 uf.union(r, col_off + v);
144 }
145 }
146
147 use std::collections::BTreeMap;
149 let mut buckets: BTreeMap<usize, (Vec<usize>, Vec<usize>)> = BTreeMap::new();
150 for r in 0..n_rows {
151 if part.row_part[r] != DMPart::Square {
152 continue;
153 }
154 let root = uf.find(r);
155 buckets.entry(root).or_default().0.push(r);
156 }
157 for v in 0..n_vars {
158 if part.col_part[v] != DMPart::Square {
159 continue;
160 }
161 let root = uf.find(col_off + v);
162 buckets.entry(root).or_default().1.push(v);
163 }
164
165 let mut comps: Vec<SquareComponent> = buckets
168 .into_values()
169 .map(|(mut rows, mut cols)| {
170 rows.sort_unstable();
171 cols.sort_unstable();
172 SquareComponent {
173 eq_rows: rows,
174 cols,
175 }
176 })
177 .filter(|c| !c.eq_rows.is_empty() || !c.cols.is_empty())
178 .collect();
179 comps.sort_by_key(|c| c.eq_rows.first().copied().unwrap_or(usize::MAX));
180
181 Self { components: comps }
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188 use crate::matching::hopcroft_karp;
189
190 fn eq_inc(n_vars: usize, n_rows: usize, edges: &[(usize, usize)]) -> EqualityIncidence {
191 let mut per_row: Vec<Vec<usize>> = vec![Vec::new(); n_rows];
192 for &(r, v) in edges {
193 per_row[r].push(v);
194 }
195 let mut adj_ptr = Vec::with_capacity(n_rows + 1);
196 let mut vars = Vec::new();
197 adj_ptr.push(0);
198 for row in per_row.iter_mut() {
199 row.sort_unstable();
200 row.dedup();
201 vars.extend_from_slice(row);
202 adj_ptr.push(vars.len());
203 }
204 EqualityIncidence {
205 n_vars,
206 eq_row_inner_idx: (0..n_rows).collect(),
207 adj_ptr,
208 vars,
209 }
210 }
211
212 fn decompose(n_vars: usize, n_rows: usize, edges: &[(usize, usize)]) -> SquareComponents {
213 let inc = eq_inc(n_vars, n_rows, edges);
214 let m = hopcroft_karp(&inc);
215 let dm = DulmageMendelsohnPartition::from_matching(&inc, &m);
216 SquareComponents::of_square_part(&inc, &m, &dm)
217 }
218
219 #[test]
220 fn components_empty_square() {
221 let c = decompose(0, 0, &[]);
222 assert!(c.components.is_empty());
223 }
224
225 #[test]
226 fn components_disjoint_3x3() {
227 let c = decompose(3, 3, &[(0, 0), (1, 1), (2, 2)]);
228 assert_eq!(c.components.len(), 3);
229 for (i, comp) in c.components.iter().enumerate() {
230 assert_eq!(comp.eq_rows, vec![i]);
231 assert_eq!(comp.cols, vec![i]);
232 }
233 }
234
235 #[test]
236 fn components_two_blocks_5x5() {
237 let edges = [
240 (0, 0),
241 (0, 1),
242 (1, 1),
243 (1, 2),
244 (2, 0),
245 (2, 2),
246 (3, 3),
247 (3, 4),
248 (4, 4),
249 ];
250 let c = decompose(5, 5, &edges);
251 assert_eq!(c.components.len(), 2);
252 assert_eq!(c.components[0].eq_rows, vec![0, 1, 2]);
253 assert_eq!(c.components[0].cols, vec![0, 1, 2]);
254 assert_eq!(c.components[1].eq_rows, vec![3, 4]);
255 assert_eq!(c.components[1].cols, vec![3, 4]);
256 }
257
258 #[test]
259 fn components_star_inside_square() {
260 let c = decompose(3, 3, &[(0, 0), (0, 1), (0, 2), (1, 1), (2, 2)]);
263 assert_eq!(c.components.len(), 1);
264 let only = &c.components[0];
265 assert_eq!(only.eq_rows, vec![0, 1, 2]);
266 assert_eq!(only.cols, vec![0, 1, 2]);
267 }
268
269 #[test]
270 fn components_order_is_deterministic() {
271 let c = decompose(
274 4,
275 4,
276 &[
277 (2, 2),
278 (2, 3),
279 (3, 2),
280 (3, 3),
281 (0, 0),
282 (0, 1),
283 (1, 0),
284 (1, 1),
285 ],
286 );
287 assert_eq!(c.components.len(), 2);
288 assert_eq!(c.components[0].eq_rows.first(), Some(&0));
290 assert_eq!(c.components[1].eq_rows.first(), Some(&2));
291 }
292
293 #[test]
294 fn components_skip_overdetermined_and_underdetermined() {
295 let c = decompose(2, 3, &[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]);
298 assert!(c.components.is_empty());
299 }
300
301 fn reference_components(
305 inc: &EqualityIncidence,
306 dm: &DulmageMendelsohnPartition,
307 ) -> Vec<(Vec<usize>, Vec<usize>)> {
308 let n_rows = inc.n_eq_rows();
309 let n_vars = inc.n_vars;
310
311 let mut col_to_rows: Vec<Vec<usize>> = vec![Vec::new(); n_vars];
313 for r in 0..n_rows {
314 for &v in inc.neighbors(r) {
315 col_to_rows[v].push(r);
316 }
317 }
318
319 let mut row_seen = vec![false; n_rows];
320 let mut col_seen = vec![false; n_vars];
321 let mut comps: Vec<(Vec<usize>, Vec<usize>)> = Vec::new();
322 for start in 0..n_rows {
323 if row_seen[start] || dm.row_part[start] != DMPart::Square {
324 continue;
325 }
326 let mut comp_rows = Vec::new();
327 let mut comp_cols = Vec::new();
328 let mut row_q: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
329 let mut col_q: std::collections::VecDeque<usize> = std::collections::VecDeque::new();
330 row_seen[start] = true;
331 row_q.push_back(start);
332 comp_rows.push(start);
333 while !row_q.is_empty() || !col_q.is_empty() {
334 while let Some(r) = row_q.pop_front() {
335 for &v in inc.neighbors(r) {
336 if dm.col_part[v] != DMPart::Square || col_seen[v] {
337 continue;
338 }
339 col_seen[v] = true;
340 comp_cols.push(v);
341 col_q.push_back(v);
342 }
343 }
344 while let Some(v) = col_q.pop_front() {
345 for &r2 in &col_to_rows[v] {
346 if dm.row_part[r2] != DMPart::Square || row_seen[r2] {
347 continue;
348 }
349 row_seen[r2] = true;
350 comp_rows.push(r2);
351 row_q.push_back(r2);
352 }
353 }
354 }
355 comp_rows.sort_unstable();
356 comp_cols.sort_unstable();
357 comps.push((comp_rows, comp_cols));
358 }
359 for v in 0..n_vars {
364 if dm.col_part[v] == DMPart::Square && !col_seen[v] {
365 comps.push((Vec::new(), vec![v]));
366 }
367 }
368 comps.sort_by_key(|(rows, cols)| {
369 rows.first()
370 .copied()
371 .unwrap_or_else(|| cols.first().copied().unwrap_or(usize::MAX))
372 });
373 comps
374 }
375
376 #[test]
383 fn components_fuzz_invariants() {
384 let mut state: u64 = 0xc0de_face_beef_b00b;
385 let mut next = || -> u64 {
386 state = state
387 .wrapping_mul(6364136223846793005)
388 .wrapping_add(1442695040888963407);
389 state >> 32
390 };
391
392 for trial in 0..30 {
393 let n_rows = 1 + (next() % 4) as usize;
394 let n_vars = 1 + (next() % 4) as usize;
395 let max_edges = (n_rows * n_vars).min(8);
396 let n_edges = (next() % (max_edges as u64 + 1)) as usize;
397
398 let mut edge_set = std::collections::BTreeSet::<(usize, usize)>::new();
399 let mut draws = 0usize;
400 while edge_set.len() < n_edges {
401 let r = (next() % n_rows as u64) as usize;
402 let v = (next() % n_vars as u64) as usize;
403 edge_set.insert((r, v));
404 draws += 1;
405 assert!(draws < 10_000);
406 }
407 let edges: Vec<(usize, usize)> = edge_set.into_iter().collect();
408
409 let inc = eq_inc(n_vars, n_rows, &edges);
410 let m = hopcroft_karp(&inc);
411 let dm = DulmageMendelsohnPartition::from_matching(&inc, &m);
412 let ours = SquareComponents::of_square_part(&inc, &m, &dm);
413 let theirs = reference_components(&inc, &dm);
414
415 assert_eq!(
417 ours.components.len(),
418 theirs.len(),
419 "trial {trial}: component count differs (ours={}, ref={})",
420 ours.components.len(),
421 theirs.len()
422 );
423 for (i, (ours_c, theirs_c)) in ours.components.iter().zip(theirs.iter()).enumerate() {
424 assert_eq!(
425 ours_c.eq_rows, theirs_c.0,
426 "trial {trial} comp {i}: rows differ"
427 );
428 assert_eq!(
429 ours_c.cols, theirs_c.1,
430 "trial {trial} comp {i}: cols differ"
431 );
432 }
433
434 let sum_rows: usize = ours.components.iter().map(|c| c.eq_rows.len()).sum();
436 let sum_cols: usize = ours.components.iter().map(|c| c.cols.len()).sum();
437 assert_eq!(sum_rows, dm.square_rows.len(), "trial {trial}");
438 assert_eq!(sum_cols, dm.square_cols.len(), "trial {trial}");
439
440 let mut col_to_comp: std::collections::HashMap<usize, usize> =
443 std::collections::HashMap::new();
444 for (i, c) in ours.components.iter().enumerate() {
445 for &v in &c.cols {
446 col_to_comp.insert(v, i);
447 }
448 }
449 for (i, c) in ours.components.iter().enumerate() {
450 for &r in &c.eq_rows {
451 for &v in inc.neighbors(r) {
452 if dm.col_part[v] != DMPart::Square {
453 continue;
454 }
455 let owner = col_to_comp.get(&v).copied().unwrap_or(usize::MAX);
456 assert_eq!(
457 owner, i,
458 "trial {trial}: edge ({r},{v}) crosses comp {i}→{owner}"
459 );
460 }
461 }
462 }
463 }
464 }
465}