toolbox_rs/
edmonds_karp.rs

1use crate::{
2    dfs::DFS,
3    edge::InputEdge,
4    graph::{Graph, NodeID},
5    max_flow::{MaxFlow, ResidualEdgeData},
6    static_graph::StaticGraph,
7};
8use bitvec::vec::BitVec;
9use itertools::Itertools;
10use log::{debug, warn};
11use std::{
12    sync::{
13        atomic::{AtomicI32, Ordering},
14        Arc,
15    },
16    time::Instant,
17};
18
19pub struct EdmondsKarp {
20    residual_graph: StaticGraph<ResidualEdgeData>,
21    max_flow: i32,
22    finished: bool,
23    source: NodeID,
24    target: NodeID,
25    bound: Option<Arc<AtomicI32>>,
26}
27
28impl MaxFlow for EdmondsKarp {
29    fn from_edge_list(
30        mut edge_list: Vec<InputEdge<ResidualEdgeData>>,
31        source: usize,
32        target: usize,
33    ) -> Self {
34        let number_of_edges = edge_list.len();
35
36        debug!("extending {} edges", edge_list.len());
37        // blindly generate reverse edges for all edges with zero capacity
38        edge_list.extend_from_within(..);
39        edge_list.iter_mut().skip(number_of_edges).for_each(|edge| {
40            edge.reverse();
41            edge.data.capacity = 0;
42        });
43        debug!("into {} edges", edge_list.len());
44
45        // dedup-merge edge set, by using the following trick: not the dedup(.) call
46        // below takes the second argument as mut. When deduping equivalent values
47        // a and b, then a is accumulated onto b.
48        edge_list.sort_unstable();
49        edge_list.dedup_by(|a, b| {
50            // edges a and b are assumed to be equivalent in the residual graph if
51            // (and only if) they are parallel. In other words, this removes parallel
52            // edges in the residual graph and accumulates capacities on the remaining
53            // egde.
54            let result = a.source == b.source && a.target == b.target;
55            if result {
56                b.data.capacity += a.data.capacity;
57            }
58            result
59        });
60        debug!("dedup-merged {} edges", edge_list.len());
61
62        // at this point the edge set of the residual graph doesn't have any
63        // duplicates anymore. note that this is fine, as we are looking to
64        // compute a node partition.
65        Self {
66            residual_graph: StaticGraph::new(edge_list),
67            max_flow: 0,
68            finished: false,
69            source,
70            target,
71            bound: None,
72        }
73    }
74    fn run_with_upper_bound(&mut self, bound: Arc<AtomicI32>) {
75        warn!("Upper bound {} is discarded", bound.load(Ordering::Relaxed));
76        self.bound = Some(bound);
77        self.run()
78    }
79
80    fn run(&mut self) {
81        let mut dfs = DFS::new(
82            &[self.source],
83            &[self.target],
84            self.residual_graph.number_of_nodes(),
85        );
86        let filter = |graph: &StaticGraph<ResidualEdgeData>, edge| graph.data(edge).capacity <= 0;
87        // let mut iteration = 0;
88        while dfs.run_with_filter(&self.residual_graph, filter) {
89            let start = Instant::now();
90            // retrieve node path. The path is unambiguous, as we removed all duplicate edges
91            // find min capacity on edges of the path
92            let bootleneck_head_tail = dfs
93                .path_iter()
94                .tuple_windows()
95                .min_by_key(|(a, b)| {
96                    let edge = self.residual_graph.find_edge_unchecked(*b, *a);
97                    self.residual_graph.data(edge).capacity
98                })
99                .expect("graph is broken, couldn't find min edge");
100            let duration = start.elapsed();
101            debug!(" flow assignment1 took: {:?} (done)", duration);
102
103            let bottleneck_edge = self
104                .residual_graph
105                .find_edge_unchecked(bootleneck_head_tail.1, bootleneck_head_tail.0);
106            debug!("  bottleneck edge: {}", bottleneck_edge);
107            let path_flow = self.residual_graph.data(bottleneck_edge).capacity;
108            debug_assert!(path_flow > 0);
109            debug!("min edge: {}, capacity: {}", bottleneck_edge, path_flow);
110            // sum up flow
111            self.max_flow += path_flow;
112            let duration = start.elapsed();
113            debug!(" flow assignment2 took: {:?} (done)", duration);
114
115            // assign flow to residual graph
116            for (a, b) in dfs.path_iter().tuple_windows() {
117                let rev_edge = self.residual_graph.find_edge_unchecked(a, b);
118                let fwd_edge = self.residual_graph.find_edge_unchecked(b, a);
119
120                self.residual_graph.data_mut(fwd_edge).capacity -= path_flow;
121                self.residual_graph.data_mut(rev_edge).capacity += path_flow;
122            }
123            let duration = start.elapsed();
124            debug!(" flow assignment3 took: {:?} (done)", duration);
125        }
126
127        self.finished = true;
128    }
129
130    fn max_flow(&self) -> Result<i32, String> {
131        if !self.finished {
132            return Err("Assigment was not computed.".to_string());
133        }
134        Ok(self.max_flow)
135    }
136
137    fn assignment(&self, source: NodeID) -> Result<BitVec, String> {
138        if !self.finished {
139            return Err("Assigment was not computed.".to_string());
140        }
141
142        // run a reachability analysis
143        let mut reachable = BitVec::with_capacity(self.residual_graph.number_of_nodes());
144        reachable.resize(self.residual_graph.number_of_nodes(), false);
145        let mut stack = vec![source];
146        stack.reserve(self.residual_graph.number_of_nodes());
147        reachable.set(source, true);
148        while let Some(node) = stack.pop() {
149            for edge in self.residual_graph.edge_range(node) {
150                let target = self.residual_graph.target(edge);
151                let reached = reachable.get(target).unwrap();
152                if !reached && self.residual_graph.data(edge).capacity > 0 {
153                    stack.push(target);
154                    reachable.set(target, true);
155                }
156            }
157        }
158        Ok(reachable)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164
165    use std::sync::atomic::AtomicI32;
166    use std::sync::Arc;
167
168    use crate::edge::InputEdge;
169    use crate::edmonds_karp::EdmondsKarp;
170    use crate::max_flow::MaxFlow;
171    use crate::max_flow::ResidualEdgeData;
172    use bitvec::bits;
173    use bitvec::prelude::Lsb0;
174
175    #[test]
176    fn max_flow_clr() {
177        let edges = vec![
178            InputEdge::new(0, 1, ResidualEdgeData::new(16)),
179            InputEdge::new(0, 2, ResidualEdgeData::new(13)),
180            InputEdge::new(1, 2, ResidualEdgeData::new(10)),
181            InputEdge::new(1, 3, ResidualEdgeData::new(12)),
182            InputEdge::new(2, 1, ResidualEdgeData::new(4)),
183            InputEdge::new(2, 4, ResidualEdgeData::new(14)),
184            InputEdge::new(3, 2, ResidualEdgeData::new(9)),
185            InputEdge::new(3, 5, ResidualEdgeData::new(20)),
186            InputEdge::new(4, 3, ResidualEdgeData::new(7)),
187            InputEdge::new(4, 5, ResidualEdgeData::new(4)),
188        ];
189
190        let source = 0;
191        let target = 5;
192        let mut max_flow_solver = EdmondsKarp::from_edge_list(edges, source, target);
193        max_flow_solver.run();
194
195        // it's OK to expect the solver to have run
196        let max_flow = max_flow_solver
197            .max_flow()
198            .expect("max flow computation did not run");
199        assert_eq!(23, max_flow);
200
201        // it's OK to expect the solver to have run
202        let assignment = max_flow_solver
203            .assignment(source)
204            .expect("assignment computation did not run");
205
206        assert_eq!(assignment, bits![1, 1, 1, 0, 1, 0]);
207    }
208
209    #[test]
210    fn run_with_upper_bound_no_effect() {
211        let edges = vec![
212            InputEdge::new(0, 1, ResidualEdgeData::new(16)),
213            InputEdge::new(0, 2, ResidualEdgeData::new(13)),
214            InputEdge::new(1, 2, ResidualEdgeData::new(10)),
215            InputEdge::new(1, 3, ResidualEdgeData::new(12)),
216            InputEdge::new(2, 1, ResidualEdgeData::new(4)),
217            InputEdge::new(2, 4, ResidualEdgeData::new(14)),
218            InputEdge::new(3, 2, ResidualEdgeData::new(9)),
219            InputEdge::new(3, 5, ResidualEdgeData::new(20)),
220            InputEdge::new(4, 3, ResidualEdgeData::new(7)),
221            InputEdge::new(4, 5, ResidualEdgeData::new(4)),
222        ];
223
224        let source = 0;
225        let target = 5;
226        let mut max_flow_solver = EdmondsKarp::from_edge_list(edges, source, target);
227        let upper_bound = Arc::new(AtomicI32::new(1));
228        max_flow_solver.run_with_upper_bound(upper_bound);
229
230        // it's OK to expect the solver to have run
231        let max_flow = max_flow_solver
232            .max_flow()
233            .expect("max flow computation did not run");
234        assert_eq!(23, max_flow);
235    }
236
237    #[test]
238    fn max_flow_ita() {
239        let edges = vec![
240            InputEdge::new(0, 1, ResidualEdgeData::new(5)),
241            InputEdge::new(0, 4, ResidualEdgeData::new(7)),
242            InputEdge::new(0, 5, ResidualEdgeData::new(6)),
243            InputEdge::new(1, 2, ResidualEdgeData::new(4)),
244            InputEdge::new(1, 7, ResidualEdgeData::new(3)),
245            InputEdge::new(4, 7, ResidualEdgeData::new(4)),
246            InputEdge::new(4, 6, ResidualEdgeData::new(1)),
247            InputEdge::new(5, 6, ResidualEdgeData::new(5)),
248            InputEdge::new(2, 3, ResidualEdgeData::new(3)),
249            InputEdge::new(7, 3, ResidualEdgeData::new(7)),
250            InputEdge::new(6, 7, ResidualEdgeData::new(1)),
251            InputEdge::new(6, 3, ResidualEdgeData::new(6)),
252        ];
253
254        let source = 0;
255        let target = 3;
256        let mut max_flow_solver = EdmondsKarp::from_edge_list(edges, source, target);
257        max_flow_solver.run();
258
259        // it's OK to expect the solver to have run
260        let max_flow = max_flow_solver
261            .max_flow()
262            .expect("max flow computation did not run");
263        assert_eq!(15, max_flow);
264
265        // it's OK to expect the solver to have run
266        let assignment = max_flow_solver
267            .assignment(source)
268            .expect("assignment computation did not run");
269        assert_eq!(assignment, bits![1, 0, 0, 0, 1, 1, 0, 0]);
270    }
271
272    #[test]
273    fn max_flow_yt() {
274        let edges = vec![
275            InputEdge::new(9, 0, ResidualEdgeData::new(5)),
276            InputEdge::new(9, 1, ResidualEdgeData::new(10)),
277            InputEdge::new(9, 2, ResidualEdgeData::new(15)),
278            InputEdge::new(0, 3, ResidualEdgeData::new(10)),
279            InputEdge::new(1, 0, ResidualEdgeData::new(15)),
280            InputEdge::new(1, 4, ResidualEdgeData::new(20)),
281            InputEdge::new(2, 5, ResidualEdgeData::new(25)),
282            InputEdge::new(3, 4, ResidualEdgeData::new(25)),
283            InputEdge::new(3, 6, ResidualEdgeData::new(10)),
284            InputEdge::new(4, 2, ResidualEdgeData::new(5)),
285            InputEdge::new(4, 7, ResidualEdgeData::new(30)),
286            InputEdge::new(5, 7, ResidualEdgeData::new(20)),
287            InputEdge::new(5, 8, ResidualEdgeData::new(10)),
288            InputEdge::new(7, 8, ResidualEdgeData::new(15)),
289            InputEdge::new(6, 10, ResidualEdgeData::new(5)),
290            InputEdge::new(7, 10, ResidualEdgeData::new(15)),
291            InputEdge::new(8, 10, ResidualEdgeData::new(10)),
292        ];
293
294        let source = 9;
295        let target = 10;
296        let mut max_flow_solver = EdmondsKarp::from_edge_list(edges, source, target);
297        max_flow_solver.run();
298
299        // it's OK to expect the solver to have run
300        let max_flow = max_flow_solver
301            .max_flow()
302            .expect("max flow computation did not run");
303        assert_eq!(30, max_flow);
304
305        // it's OK to expect the solver to have run
306        let assignment = max_flow_solver
307            .assignment(source)
308            .expect("assignment computation did not run");
309        assert_eq!(assignment, bits![0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0]);
310    }
311
312    #[test]
313    fn max_flow_ff() {
314        let edges = vec![
315            InputEdge::new(0, 1, ResidualEdgeData::new(7)),
316            InputEdge::new(0, 2, ResidualEdgeData::new(3)),
317            InputEdge::new(1, 2, ResidualEdgeData::new(1)),
318            InputEdge::new(1, 3, ResidualEdgeData::new(6)),
319            InputEdge::new(2, 4, ResidualEdgeData::new(8)),
320            InputEdge::new(3, 5, ResidualEdgeData::new(2)),
321            InputEdge::new(3, 2, ResidualEdgeData::new(3)),
322            InputEdge::new(4, 3, ResidualEdgeData::new(2)),
323            InputEdge::new(4, 5, ResidualEdgeData::new(8)),
324        ];
325
326        let source = 0;
327        let target = 5;
328        let mut max_flow_solver = EdmondsKarp::from_edge_list(edges, source, target);
329        max_flow_solver.run();
330
331        // it's OK to expect the solver to have run
332        let max_flow = max_flow_solver
333            .max_flow()
334            .expect("max flow computation did not run");
335        assert_eq!(9, max_flow);
336
337        // it's OK to expect the solver to have run
338        let assignment = max_flow_solver
339            .assignment(source)
340            .expect("assignment computation did not run");
341        assert_eq!(assignment, bits![1, 1, 0, 1, 0, 0]);
342    }
343
344    #[test]
345    #[should_panic]
346    fn flow_not_computed() {
347        let edges = vec![
348            InputEdge::new(0, 1, ResidualEdgeData::new(7)),
349            InputEdge::new(0, 2, ResidualEdgeData::new(3)),
350            InputEdge::new(1, 2, ResidualEdgeData::new(1)),
351            InputEdge::new(1, 3, ResidualEdgeData::new(6)),
352            InputEdge::new(2, 4, ResidualEdgeData::new(8)),
353            InputEdge::new(3, 5, ResidualEdgeData::new(2)),
354            InputEdge::new(3, 2, ResidualEdgeData::new(3)),
355            InputEdge::new(4, 3, ResidualEdgeData::new(2)),
356            InputEdge::new(4, 5, ResidualEdgeData::new(8)),
357        ];
358
359        // the expect(.) call is being tested
360        EdmondsKarp::from_edge_list(edges, 0, 1)
361            .max_flow()
362            .expect("max flow computation did not run");
363    }
364
365    #[test]
366    #[should_panic]
367    fn assignment_not_computed() {
368        let edges = vec![
369            InputEdge::new(0, 1, ResidualEdgeData::new(7)),
370            InputEdge::new(0, 2, ResidualEdgeData::new(3)),
371            InputEdge::new(1, 2, ResidualEdgeData::new(1)),
372            InputEdge::new(1, 3, ResidualEdgeData::new(6)),
373            InputEdge::new(2, 4, ResidualEdgeData::new(8)),
374            InputEdge::new(3, 5, ResidualEdgeData::new(2)),
375            InputEdge::new(3, 2, ResidualEdgeData::new(3)),
376            InputEdge::new(4, 3, ResidualEdgeData::new(2)),
377            InputEdge::new(4, 5, ResidualEdgeData::new(8)),
378        ];
379
380        // the expect(.) call is being tested
381        EdmondsKarp::from_edge_list(edges, 0, 1)
382            .assignment(0)
383            .expect("assignment computation did not run");
384    }
385}