1use crate::graph::traits::{GraphBase, GraphQuery};
15use crate::graph::Graph;
16use crate::node::NodeIndex;
17use std::collections::{HashMap, VecDeque};
18
19pub fn edmonds_karp<T, E, F>(
59 graph: &mut Graph<T, E>,
60 source: NodeIndex,
61 sink: NodeIndex,
62 mut capacity: F,
63) -> f64
64where
65 F: FnMut(NodeIndex, NodeIndex, &E) -> f64,
66{
67 let n = graph.node_count();
68 if n == 0 {
69 return 0.0;
70 }
71
72 let node_indices: Vec<NodeIndex> = graph.nodes().map(|n| n.index()).collect();
74 let index_to_pos: HashMap<usize, usize> = node_indices
75 .iter()
76 .enumerate()
77 .map(|(i, ni)| (ni.index(), i))
78 .collect();
79
80 let mut adj: HashMap<usize, Vec<(usize, f64)>> = HashMap::with_capacity(n);
84
85 for edge in graph.edges() {
87 if let Ok((src, tgt)) = graph.edge_endpoints(edge.index()) {
88 if let Some(&u) = index_to_pos.get(&src.index()) {
89 if let Some(&v) = index_to_pos.get(&tgt.index()) {
90 let cap = capacity(src, tgt, edge.data());
91 if cap > 0.0 {
92 adj.entry(u).or_default().push((v, cap));
93 adj.entry(v).or_default().push((u, 0.0));
95 }
96 }
97 }
98 }
99 }
100
101 let mut max_flow = 0.0;
102
103 fn bfs(
105 adj: &HashMap<usize, Vec<(usize, f64)>>,
106 source_pos: usize,
107 sink_pos: usize,
108 n: usize,
109 parent: &mut [(isize, usize)], ) -> bool {
111 parent.fill((-1, 0));
112 let mut visited = vec![false; n];
113 let mut queue = VecDeque::new();
114
115 queue.push_back(source_pos);
116 visited[source_pos] = true;
117
118 while let Some(u) = queue.pop_front() {
119 if let Some(neighbors) = adj.get(&u) {
120 for (idx, (v, cap)) in neighbors.iter().enumerate() {
121 if !visited[*v] && *cap > 1e-9 {
122 visited[*v] = true;
123 parent[*v] = (u as isize, idx);
124 if *v == sink_pos {
125 return true;
126 }
127 queue.push_back(*v);
128 }
129 }
130 }
131 }
132 false
133 }
134
135 let source_pos = *index_to_pos.get(&source.index()).unwrap_or(&0);
136 let sink_pos = *index_to_pos.get(&sink.index()).unwrap_or(&0);
137 let mut parent = vec![(-1, 0); n];
138
139 while bfs(&adj, source_pos, sink_pos, n, &mut parent) {
140 let mut path_flow = f64::INFINITY;
142 let mut v = sink_pos;
143 while v != source_pos {
144 let (u, edge_idx) = parent[v];
145 let u = u as usize;
146 if let Some(neighbors) = adj.get(&u) {
147 if let Some((_, cap)) = neighbors.get(edge_idx) {
148 path_flow = path_flow.min(*cap);
149 }
150 }
151 v = u;
152 }
153
154 let mut v = sink_pos;
156 while v != source_pos {
157 let (u, edge_idx) = parent[v];
158 let u = u as usize;
159
160 if let Some(neighbors) = adj.get_mut(&u) {
162 if let Some((_target, cap)) = neighbors.get_mut(edge_idx) {
163 *cap -= path_flow;
164 let target_val = *_target;
165 if let Some(reverse_neighbors) = adj.get_mut(&v) {
167 if let Some((_, rev_cap)) =
168 reverse_neighbors.iter_mut().find(|(t, _)| *t == target_val)
169 {
170 *rev_cap += path_flow;
171 }
172 }
173 }
174 }
175 v = u;
176 }
177
178 max_flow += path_flow;
179 }
180
181 max_flow
182}
183
184pub fn dinic<T, E, F>(
197 graph: &mut Graph<T, E>,
198 source: NodeIndex,
199 sink: NodeIndex,
200 mut capacity: F,
201) -> f64
202where
203 F: FnMut(NodeIndex, NodeIndex, &E) -> f64,
204{
205 let n = graph.node_count();
206 if n == 0 {
207 return 0.0;
208 }
209
210 let node_indices: Vec<NodeIndex> = graph.nodes().map(|n| n.index()).collect();
211 let index_to_pos: HashMap<usize, usize> = node_indices
212 .iter()
213 .enumerate()
214 .map(|(i, ni)| (ni.index(), i))
215 .collect();
216
217 let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
220 let mut residual: HashMap<usize, HashMap<usize, f64>> = HashMap::with_capacity(n);
221
222 for edge in graph.edges() {
223 if let Ok((src, tgt)) = graph.edge_endpoints(edge.index()) {
224 if let (Some(&u), Some(&v)) = (
225 index_to_pos.get(&src.index()),
226 index_to_pos.get(&tgt.index()),
227 ) {
228 let cap = capacity(src, tgt, edge.data());
229 if !residual
230 .get(&u)
231 .map(|m| m.contains_key(&v))
232 .unwrap_or(false)
233 {
234 adj[u].push(v);
235 }
236 residual.entry(u).or_default().insert(v, cap);
237 }
238 }
239 }
240
241 let source_pos = *index_to_pos.get(&source.index()).unwrap_or(&0);
242 let sink_pos = *index_to_pos.get(&sink.index()).unwrap_or(&0);
243 let mut max_flow = 0.0;
244
245 fn bfs_level(
247 adj: &[Vec<usize>],
248 residual: &HashMap<usize, HashMap<usize, f64>>,
249 source_pos: usize,
250 sink_pos: usize,
251 level: &mut [isize],
252 ) -> bool {
253 let _n = adj.len();
254 level.fill(-1);
255 level[source_pos] = 0;
256 let mut queue = VecDeque::new();
257 queue.push_back(source_pos);
258
259 while let Some(u) = queue.pop_front() {
260 if let Some(neighbors) = adj.get(u) {
261 for &v in neighbors {
262 let cap = residual
263 .get(&u)
264 .and_then(|m| m.get(&v))
265 .copied()
266 .unwrap_or(0.0);
267 if level[v] == -1 && cap > 1e-9 {
268 level[v] = level[u] + 1;
269 if v == sink_pos {
270 return true;
271 }
272 queue.push_back(v);
273 }
274 }
275 }
276 }
277 false
278 }
279
280 fn dfs_block(
282 u: usize,
283 flow: f64,
284 sink_pos: usize,
285 adj: &[Vec<usize>],
286 residual: &mut HashMap<usize, HashMap<usize, f64>>,
287 level: &mut [isize],
288 ptr: &mut [usize],
289 ) -> f64 {
290 if u == sink_pos || flow < 1e-9 {
291 return flow;
292 }
293
294 let start = ptr[u];
295 for i in start..adj[u].len() {
296 ptr[u] = i;
297 let v = adj[u][i];
298 let cap = residual
299 .get(&u)
300 .and_then(|m| m.get(&v))
301 .copied()
302 .unwrap_or(0.0);
303 if level[v] == level[u] + 1 && cap > 1e-9 {
304 let pushed = dfs_block(v, flow.min(cap), sink_pos, adj, residual, level, ptr);
305 if pushed > 1e-9 {
306 *residual.entry(u).or_default().get_mut(&v).unwrap() -= pushed;
307 *residual.entry(v).or_default().entry(u).or_insert(0.0) += pushed;
308 return pushed;
309 }
310 }
311 }
312 0.0
313 }
314
315 let mut level = vec![-1; n];
316 let mut ptr = vec![0; n];
317
318 while bfs_level(&adj, &residual, source_pos, sink_pos, &mut level) {
319 ptr.fill(0);
320 while let Some(push) = {
321 let pushed = dfs_block(
322 source_pos,
323 f64::INFINITY,
324 sink_pos,
325 &adj,
326 &mut residual,
327 &mut level,
328 &mut ptr,
329 );
330 if pushed > 1e-9 {
331 Some(pushed)
332 } else {
333 None
334 }
335 } {
336 max_flow += push;
337 }
338 }
339
340 max_flow
341}
342
343pub fn push_relabel<T, E, F>(
387 graph: &mut Graph<T, E>,
388 source: NodeIndex,
389 sink: NodeIndex,
390 mut capacity: F,
391) -> f64
392where
393 F: FnMut(NodeIndex, NodeIndex, &E) -> f64,
394{
395 let n = graph.node_count();
396 if n == 0 {
397 return 0.0;
398 }
399
400 let node_indices: Vec<NodeIndex> = graph.nodes().map(|n| n.index()).collect();
402 let index_to_pos: HashMap<usize, usize> = node_indices
403 .iter()
404 .enumerate()
405 .map(|(i, ni)| (ni.index(), i))
406 .collect();
407
408 let source_pos = index_to_pos[&source.index()];
409 let sink_pos = index_to_pos[&sink.index()];
410
411 let mut adj: Vec<Vec<usize>> = vec![vec![]; n];
414 let mut residual: HashMap<usize, HashMap<usize, f64>> = HashMap::with_capacity(n);
415
416 for edge in graph.edges() {
418 if let Ok((src, tgt)) = graph.edge_endpoints(edge.index()) {
419 if let Some(&u) = index_to_pos.get(&src.index()) {
420 if let Some(&v) = index_to_pos.get(&tgt.index()) {
421 let cap = capacity(src, tgt, edge.data());
422 residual.entry(u).or_default().insert(v, cap);
423 adj[u].push(v);
424 adj[v].push(u); }
426 }
427 }
428 }
429
430 let mut height = vec![0; n];
432 let mut excess = vec![0.0; n];
433
434 height[source_pos] = n;
436
437 for &v in &adj[source_pos] {
438 if let Some(cap) = residual.get(&source_pos).and_then(|m| m.get(&v)).copied() {
439 if cap > 1e-9 {
440 *residual
441 .entry(source_pos)
442 .or_default()
443 .entry(v)
444 .or_insert(0.0) = 0.0;
445 *residual
446 .entry(v)
447 .or_default()
448 .entry(source_pos)
449 .or_insert(0.0) = cap;
450 excess[v] = cap;
451 excess[source_pos] -= cap;
452 }
453 }
454 }
455
456 let mut queue: VecDeque<usize> = VecDeque::new();
458 let mut in_queue = vec![false; n];
459
460 for i in 0..n {
461 if i != source_pos && i != sink_pos && excess[i] > 1e-9 {
462 queue.push_back(i);
463 in_queue[i] = true;
464 }
465 }
466
467 while let Some(u) = queue.pop_front() {
469 in_queue[u] = false;
470
471 if excess[u] <= 1e-9 {
472 continue;
473 }
474
475 let mut pushed = false;
477 if let Some(neighbors) = adj.get(u) {
478 for &v in neighbors {
479 let cap = residual
480 .get(&u)
481 .and_then(|m| m.get(&v))
482 .copied()
483 .unwrap_or(0.0);
484 if cap > 1e-9 && height[u] == height[v] + 1 {
485 let delta = excess[u].min(cap);
486 *residual.entry(u).or_default().entry(v).or_insert(0.0) -= delta;
487 *residual.entry(v).or_default().entry(u).or_insert(0.0) += delta;
488 excess[u] -= delta;
489 excess[v] += delta;
490
491 if v != source_pos && v != sink_pos && !in_queue[v] && excess[v] > 1e-9 {
492 queue.push_back(v);
493 in_queue[v] = true;
494 }
495
496 if excess[u] <= 1e-9 {
497 pushed = true;
498 break;
499 }
500 }
501 }
502 }
503
504 if !pushed && excess[u] > 1e-9 {
506 let mut min_height = usize::MAX;
507 if let Some(neighbors) = adj.get(u) {
508 for &v in neighbors {
509 let cap = residual
510 .get(&u)
511 .and_then(|m| m.get(&v))
512 .copied()
513 .unwrap_or(0.0);
514 if cap > 1e-9 && height[v] < min_height {
515 min_height = height[v];
516 }
517 }
518 }
519
520 if min_height < usize::MAX {
521 height[u] = min_height + 1;
522 }
523
524 if !in_queue[u] {
526 queue.push_back(u);
527 in_queue[u] = true;
528 }
529 }
530 }
531
532 let mut max_flow = 0.0;
534 if let Some(neighbors) = adj.get(source_pos) {
535 for &v in neighbors {
536 let cap = residual
537 .get(&v)
538 .and_then(|m| m.get(&source_pos))
539 .copied()
540 .unwrap_or(0.0);
541 if cap > 1e-9 {
542 max_flow += cap;
543 }
544 }
545 }
546
547 max_flow
548}
549#[cfg(test)]
550mod tests {
551 use super::*;
552 use crate::graph::builders::GraphBuilder;
553
554 #[test]
555 fn test_edmonds_karp_basic() {
556 let mut graph = GraphBuilder::directed()
557 .with_nodes(vec!["S", "A", "B", "C", "D", "T"])
558 .with_edges(vec![
559 (0, 1, 16.0), (0, 2, 13.0), (1, 2, 10.0), (1, 3, 12.0), (2, 1, 4.0), (2, 4, 14.0), (3, 1, 9.0), (3, 5, 20.0), (4, 3, 7.0), (4, 5, 4.0), ])
570 .build()
571 .unwrap();
572
573 let source = NodeIndex::new(0, 1);
574 let sink = NodeIndex::new(5, 1);
575 let max_flow = edmonds_karp(&mut graph, source, sink, |_, _, cap| *cap);
576
577 assert_eq!(max_flow, 23.0);
578 }
579
580 #[test]
581 fn test_dinic_basic() {
582 let mut graph = GraphBuilder::directed()
583 .with_nodes(vec!["S", "A", "B", "T"])
584 .with_edges(vec![
585 (0, 1, 10.0), (0, 2, 5.0), (1, 2, 15.0), (1, 3, 10.0), (2, 3, 10.0), ])
591 .build()
592 .unwrap();
593
594 let source = NodeIndex::new(0, 1);
595 let sink = NodeIndex::new(3, 1);
596 let max_flow = dinic(&mut graph, source, sink, |_, _, cap| *cap);
597
598 assert_eq!(max_flow, 15.0);
599 }
600
601 #[test]
602 fn test_edmonds_karp_simple() {
603 let mut graph = GraphBuilder::directed()
604 .with_nodes(vec!["S", "A", "T"])
605 .with_edges(vec![
606 (0, 1, 5.0), (1, 2, 3.0), ])
609 .build()
610 .unwrap();
611
612 let source = NodeIndex::new(0, 1);
613 let sink = NodeIndex::new(2, 1);
614 let max_flow = edmonds_karp(&mut graph, source, sink, |_, _, cap| *cap);
615
616 assert_eq!(max_flow, 3.0); }
618
619 #[test]
620 fn test_push_relabel_basic() {
621 let mut graph = GraphBuilder::directed()
622 .with_nodes(vec!["S", "A", "B", "T"])
623 .with_edges(vec![
624 (0, 1, 10.0), (0, 2, 5.0), (1, 2, 15.0), (1, 3, 10.0), (2, 3, 10.0), ])
630 .build()
631 .unwrap();
632
633 let source = NodeIndex::new(0, 1);
634 let sink = NodeIndex::new(3, 1);
635 let max_flow = push_relabel(&mut graph, source, sink, |_, _, cap| *cap);
636
637 assert_eq!(max_flow, 15.0);
638 }
639
640 #[test]
641 fn test_edmonds_karp_no_path() {
642 let mut graph = GraphBuilder::directed()
644 .with_nodes(vec!["S", "A", "B", "T"])
645 .with_edges(vec![
646 (0, 1, 10.0), (2, 3, 5.0), ])
649 .build()
650 .unwrap();
651
652 let source = NodeIndex::new(0, 1);
653 let sink = NodeIndex::new(3, 1);
654 let max_flow = edmonds_karp(&mut graph, source, sink, |_, _, cap| *cap);
655
656 assert_eq!(max_flow, 0.0);
657 }
658
659 #[test]
660 fn test_dinic_no_path() {
661 let mut graph = GraphBuilder::directed()
662 .with_nodes(vec!["S", "A", "B", "T"])
663 .with_edges(vec![
664 (0, 1, 10.0), (2, 3, 5.0), ])
667 .build()
668 .unwrap();
669
670 let source = NodeIndex::new(0, 1);
671 let sink = NodeIndex::new(3, 1);
672 let max_flow = dinic(&mut graph, source, sink, |_, _, cap| *cap);
673
674 assert_eq!(max_flow, 0.0);
675 }
676
677 #[test]
678 fn test_edmonds_karp_single_edge() {
679 let mut graph = GraphBuilder::directed()
680 .with_nodes(vec!["S", "T"])
681 .with_edges(vec![(0, 1, 42.0)])
682 .build()
683 .unwrap();
684
685 let source = NodeIndex::new(0, 1);
686 let sink = NodeIndex::new(1, 1);
687 let max_flow = edmonds_karp(&mut graph, source, sink, |_, _, cap| *cap);
688
689 assert_eq!(max_flow, 42.0);
690 }
691
692 #[test]
693 fn test_dinic_single_edge() {
694 let mut graph = GraphBuilder::directed()
695 .with_nodes(vec!["S", "T"])
696 .with_edges(vec![(0, 1, 42.0)])
697 .build()
698 .unwrap();
699
700 let source = NodeIndex::new(0, 1);
701 let sink = NodeIndex::new(1, 1);
702 let max_flow = dinic(&mut graph, source, sink, |_, _, cap| *cap);
703
704 assert_eq!(max_flow, 42.0);
705 }
706
707 #[test]
708 fn test_push_relabel_no_path() {
709 let mut graph = GraphBuilder::directed()
710 .with_nodes(vec!["S", "A", "B", "T"])
711 .with_edges(vec![
712 (0, 1, 10.0), (2, 3, 5.0), ])
715 .build()
716 .unwrap();
717
718 let source = NodeIndex::new(0, 1);
719 let sink = NodeIndex::new(3, 1);
720 let max_flow = push_relabel(&mut graph, source, sink, |_, _, cap| *cap);
721
722 assert_eq!(max_flow, 0.0);
723 }
724
725 #[test]
726 fn test_edmonds_karp_zero_capacity() {
727 let mut graph = GraphBuilder::directed()
728 .with_nodes(vec!["S", "A", "T"])
729 .with_edges(vec![
730 (0, 1, 0.0), (1, 2, 5.0), ])
733 .build()
734 .unwrap();
735
736 let source = NodeIndex::new(0, 1);
737 let sink = NodeIndex::new(2, 1);
738 let max_flow = edmonds_karp(&mut graph, source, sink, |_, _, cap| *cap);
739
740 assert_eq!(max_flow, 0.0);
741 }
742}