1use crate::graph::{VertexId, Weight};
13use std::collections::{HashMap, HashSet, VecDeque};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum EdgeColor {
19 Red,
21 Blue,
23 Green,
25 Yellow,
27}
28
29#[derive(Debug, Clone)]
32pub struct EdgeColoring {
33 colors: HashMap<(VertexId, VertexId), EdgeColor>,
35 pub a: usize,
37 pub b: usize,
39}
40
41impl EdgeColoring {
42 pub fn new(a: usize, b: usize) -> Self {
44 Self {
45 colors: HashMap::new(),
46 a,
47 b,
48 }
49 }
50
51 pub fn get(&self, u: VertexId, v: VertexId) -> Option<EdgeColor> {
53 let key = if u < v { (u, v) } else { (v, u) };
54 self.colors.get(&key).copied()
55 }
56
57 pub fn set(&mut self, u: VertexId, v: VertexId, color: EdgeColor) {
59 let key = if u < v { (u, v) } else { (v, u) };
60 self.colors.insert(key, color);
61 }
62
63 pub fn has_color(&self, u: VertexId, v: VertexId, color: EdgeColor) -> bool {
65 self.get(u, v) == Some(color)
66 }
67}
68
69pub fn generate_coloring_family(a: usize, b: usize, num_edges: usize) -> Vec<EdgeColoring> {
72 let log_n = (num_edges.max(2) as f64).log2().ceil() as usize;
76 let family_size = (1 << (a.min(b) * (a + b).max(1).ilog2() as usize + 1)) * log_n;
77 let family_size = family_size.min(100); let mut family = Vec::with_capacity(family_size);
80
81 for seed in 0..family_size {
82 let coloring = EdgeColoring::new(a, b);
83 family.push(coloring);
86 }
87
88 family
89}
90
91#[derive(Debug, Clone)]
93pub struct GreedyForestPacking {
94 pub num_forests: usize,
96 edge_forest: HashMap<(VertexId, VertexId), usize>,
98 forests: Vec<HashSet<(VertexId, VertexId)>>,
100 forest_parents: Vec<HashMap<VertexId, VertexId>>,
102}
103
104impl GreedyForestPacking {
105 pub fn new(num_forests: usize) -> Self {
108 Self {
109 num_forests,
110 edge_forest: HashMap::new(),
111 forests: vec![HashSet::new(); num_forests],
112 forest_parents: vec![HashMap::new(); num_forests],
113 }
114 }
115
116 fn find_root(&mut self, forest_id: usize, v: VertexId) -> VertexId {
118 if !self.forest_parents[forest_id].contains_key(&v) {
119 self.forest_parents[forest_id].insert(v, v);
120 return v;
121 }
122
123 let parent = self.forest_parents[forest_id][&v];
124 if parent == v {
125 return v;
126 }
127
128 let root = self.find_root(forest_id, parent);
129 self.forest_parents[forest_id].insert(v, root);
130 root
131 }
132
133 fn union(&mut self, forest_id: usize, u: VertexId, v: VertexId) {
135 let root_u = self.find_root(forest_id, u);
136 let root_v = self.find_root(forest_id, v);
137 if root_u != root_v {
138 self.forest_parents[forest_id].insert(root_u, root_v);
139 }
140 }
141
142 fn would_create_cycle(&mut self, forest_id: usize, u: VertexId, v: VertexId) -> bool {
144 self.find_root(forest_id, u) == self.find_root(forest_id, v)
145 }
146
147 pub fn insert_edge(&mut self, u: VertexId, v: VertexId) -> Option<usize> {
149 let key = if u < v { (u, v) } else { (v, u) };
150
151 if self.edge_forest.contains_key(&key) {
153 return self.edge_forest.get(&key).copied();
154 }
155
156 for forest_id in 0..self.num_forests {
158 if !self.would_create_cycle(forest_id, u, v) {
159 self.forests[forest_id].insert(key);
160 self.edge_forest.insert(key, forest_id);
161 self.union(forest_id, u, v);
162 return Some(forest_id);
163 }
164 }
165
166 None
168 }
169
170 pub fn delete_edge(&mut self, u: VertexId, v: VertexId) -> Option<usize> {
172 let key = if u < v { (u, v) } else { (v, u) };
173
174 if let Some(forest_id) = self.edge_forest.remove(&key) {
175 self.forests[forest_id].remove(&key);
176 self.rebuild_forest_connectivity(forest_id);
178 return Some(forest_id);
179 }
180 None
181 }
182
183 fn rebuild_forest_connectivity(&mut self, forest_id: usize) {
185 self.forest_parents[forest_id].clear();
186 let edges: Vec<_> = self.forests[forest_id].iter().copied().collect();
188 for (u, v) in edges {
189 self.union(forest_id, u, v);
190 }
191 }
192
193 pub fn is_tree_edge(&self, u: VertexId, v: VertexId) -> bool {
195 let key = if u < v { (u, v) } else { (v, u) };
196 self.edge_forest.contains_key(&key)
197 }
198
199 pub fn get_forest(&self, u: VertexId, v: VertexId) -> Option<usize> {
201 let key = if u < v { (u, v) } else { (v, u) };
202 self.edge_forest.get(&key).copied()
203 }
204
205 pub fn forest_edges(&self, forest_id: usize) -> &HashSet<(VertexId, VertexId)> {
207 &self.forests[forest_id]
208 }
209}
210
211#[derive(Debug, Clone)]
213pub struct LocalCut {
214 pub vertices: HashSet<VertexId>,
216 pub boundary: Vec<(VertexId, VertexId)>,
218 pub cut_value: f64,
220 pub volume: usize,
222}
223
224#[derive(Debug)]
227pub struct DeterministicLocalKCut {
228 lambda_max: u64,
230 nu: usize,
232 beta: usize,
234 forests: GreedyForestPacking,
236 red_blue_colorings: Vec<EdgeColoring>,
238 green_yellow_colorings: Vec<EdgeColoring>,
240 adjacency: HashMap<VertexId, HashMap<VertexId, Weight>>,
242 edges: HashSet<(VertexId, VertexId)>,
244}
245
246impl DeterministicLocalKCut {
247 pub fn new(lambda_max: u64, nu: usize, beta: usize) -> Self {
249 let num_forests = ((6 * lambda_max) as usize).max(10);
251
252 let a_rb = 2 * beta;
254 let b_rb = nu;
255 let a_gy = 2 * beta - 1;
256 let b_gy = lambda_max as usize;
257
258 Self {
259 lambda_max,
260 nu,
261 beta,
262 forests: GreedyForestPacking::new(num_forests),
263 red_blue_colorings: generate_coloring_family(a_rb, b_rb, 1000),
264 green_yellow_colorings: generate_coloring_family(a_gy, b_gy, 1000),
265 adjacency: HashMap::new(),
266 edges: HashSet::new(),
267 }
268 }
269
270 pub fn insert_edge(&mut self, u: VertexId, v: VertexId, weight: Weight) {
272 let key = if u < v { (u, v) } else { (v, u) };
273
274 if self.edges.contains(&key) {
275 return;
276 }
277
278 self.edges.insert(key);
279 self.adjacency.entry(u).or_default().insert(v, weight);
280 self.adjacency.entry(v).or_default().insert(u, weight);
281
282 if let Some(forest_id) = self.forests.insert_edge(u, v) {
284 for coloring in &mut self.red_blue_colorings {
286 let color = if (u + v + forest_id as u64) % 2 == 0 {
287 EdgeColor::Blue
288 } else {
289 EdgeColor::Red
290 };
291 coloring.set(u, v, color);
292 }
293 } else {
294 for coloring in &mut self.green_yellow_colorings {
296 let color = if (u * v) % 2 == 0 {
297 EdgeColor::Green
298 } else {
299 EdgeColor::Yellow
300 };
301 coloring.set(u, v, color);
302 }
303 }
304 }
305
306 pub fn delete_edge(&mut self, u: VertexId, v: VertexId) {
308 let key = if u < v { (u, v) } else { (v, u) };
309
310 if !self.edges.remove(&key) {
311 return;
312 }
313
314 if let Some(neighbors) = self.adjacency.get_mut(&u) {
315 neighbors.remove(&v);
316 }
317 if let Some(neighbors) = self.adjacency.get_mut(&v) {
318 neighbors.remove(&u);
319 }
320
321 self.forests.delete_edge(u, v);
322 }
323
324 pub fn query(&self, v: VertexId) -> Vec<LocalCut> {
327 let mut results = Vec::new();
328 let mut seen_cuts: HashSet<Vec<VertexId>> = HashSet::new();
329
330 for forest_id in 0..self.forests.num_forests {
332 for rb_coloring in &self.red_blue_colorings {
333 for gy_coloring in &self.green_yellow_colorings {
334 if let Some(cut) = self.color_coded_dfs(v, forest_id, rb_coloring, gy_coloring)
336 {
337 let mut sorted_vertices: Vec<_> = cut.vertices.iter().copied().collect();
339 sorted_vertices.sort();
340
341 if !seen_cuts.contains(&sorted_vertices)
342 && cut.cut_value <= self.lambda_max as f64
343 {
344 seen_cuts.insert(sorted_vertices);
345 results.push(cut);
346 }
347 }
348 }
349 }
350 }
351
352 results
353 }
354
355 fn color_coded_dfs(
359 &self,
360 start: VertexId,
361 _forest_id: usize,
362 rb_coloring: &EdgeColoring,
363 gy_coloring: &EdgeColoring,
364 ) -> Option<LocalCut> {
365 let mut visited = HashSet::new();
366 let mut stack = vec![start];
367 let mut volume = 0usize;
368 let mut boundary = Vec::new();
369
370 while let Some(u) = stack.pop() {
371 if visited.contains(&u) {
372 continue;
373 }
374 visited.insert(u);
375
376 if let Some(neighbors) = self.adjacency.get(&u) {
378 volume += neighbors.len();
379
380 if volume > self.nu {
381 return None;
383 }
384
385 for (&v, &_weight) in neighbors {
386 let is_tree_edge = self.forests.is_tree_edge(u, v);
387
388 if is_tree_edge {
389 if rb_coloring.has_color(u, v, EdgeColor::Blue) {
391 if !visited.contains(&v) {
392 stack.push(v);
393 }
394 } else {
395 boundary.push((u, v));
397 }
398 } else {
399 if gy_coloring.has_color(u, v, EdgeColor::Green) {
401 if !visited.contains(&v) {
402 stack.push(v);
403 }
404 } else {
405 if !visited.contains(&v) {
407 boundary.push((u, v));
408 }
409 }
410 }
411 }
412 }
413 }
414
415 let cut_value: f64 = boundary
417 .iter()
418 .map(|&(u, v)| {
419 self.adjacency
420 .get(&u)
421 .and_then(|n| n.get(&v))
422 .copied()
423 .unwrap_or(1.0)
424 })
425 .sum();
426
427 Some(LocalCut {
428 vertices: visited,
429 boundary,
430 cut_value,
431 volume,
432 })
433 }
434
435 pub fn vertices(&self) -> Vec<VertexId> {
437 self.adjacency.keys().copied().collect()
438 }
439
440 pub fn neighbors(&self, v: VertexId) -> Vec<(VertexId, Weight)> {
442 self.adjacency
443 .get(&v)
444 .map(|n| n.iter().map(|(&v, &w)| (v, w)).collect())
445 .unwrap_or_default()
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::*;
452
453 #[test]
454 fn test_forest_packing_basic() {
455 let mut packing = GreedyForestPacking::new(3);
456
457 assert!(packing.insert_edge(1, 2).is_some());
459 assert!(packing.insert_edge(2, 3).is_some());
460 assert!(packing.insert_edge(3, 4).is_some());
461
462 assert!(packing.is_tree_edge(1, 2));
463 assert!(packing.is_tree_edge(2, 3));
464 }
465
466 #[test]
467 fn test_forest_packing_cycle() {
468 let mut packing = GreedyForestPacking::new(3);
469
470 packing.insert_edge(1, 2);
472 packing.insert_edge(2, 3);
473 let forest = packing.insert_edge(1, 3);
475
476 assert!(forest.is_some());
478 }
479
480 #[test]
481 fn test_localkcut_query() {
482 let mut lkc = DeterministicLocalKCut::new(10, 100, 2);
483
484 lkc.insert_edge(1, 2, 1.0);
486 lkc.insert_edge(2, 3, 1.0);
487 lkc.insert_edge(3, 4, 1.0);
488
489 let cuts = lkc.query(1);
490
491 assert!(!cuts.is_empty());
493 assert!(cuts.iter().any(|c| c.vertices.contains(&1)));
494 }
495
496 #[test]
497 fn test_coloring_family() {
498 let family = generate_coloring_family(2, 5, 100);
499 assert!(!family.is_empty());
500 }
501
502 #[test]
503 fn test_edge_deletion() {
504 let mut lkc = DeterministicLocalKCut::new(10, 100, 2);
505
506 lkc.insert_edge(1, 2, 1.0);
507 lkc.insert_edge(2, 3, 1.0);
508
509 lkc.delete_edge(1, 2);
510
511 assert!(!lkc.forests.is_tree_edge(1, 2));
513 }
514}