competitive_programming_rs/graph/
maximum_flow.rs

1pub mod dinitz {
2    const INF: i64 = 1 << 60;
3    pub struct Edge {
4        pub to: usize,
5        pub rev: usize,
6        pub is_reversed: bool,
7        pub cap: i64,
8    }
9
10    pub struct Dinitz {
11        pub g: Vec<Vec<Edge>>,
12    }
13
14    impl Dinitz {
15        pub fn new(v: usize) -> Dinitz {
16            let mut g: Vec<Vec<Edge>> = Vec::new();
17            for _ in 0..v {
18                g.push(Vec::new());
19            }
20            Dinitz { g }
21        }
22
23        pub fn add_edge(&mut self, from: usize, to: usize, cap: i64) {
24            let to_len = self.g[to].len();
25            let from_len = self.g[from].len();
26            self.g[from].push(Edge {
27                to,
28                rev: to_len,
29                is_reversed: false,
30                cap,
31            });
32            self.g[to].push(Edge {
33                to: from,
34                rev: from_len,
35                is_reversed: true,
36                cap: 0,
37            });
38        }
39
40        fn dfs(
41            &mut self,
42            v: usize,
43            sink: usize,
44            flow: i64,
45            level: &[i32],
46            iter: &mut [usize],
47        ) -> i64 {
48            if v == sink {
49                return flow;
50            }
51            while iter[v] < self.g[v].len() {
52                let flow = std::cmp::min(flow, self.g[v][iter[v]].cap);
53                let to = self.g[v][iter[v]].to;
54                if flow > 0 && level[v] < level[to] {
55                    let flowed = self.dfs(to, sink, flow, level, iter);
56                    if flowed > 0 {
57                        let rev = self.g[v][iter[v]].rev;
58                        self.g[v][iter[v]].cap -= flowed;
59                        self.g[to][rev].cap += flowed;
60                        return flowed;
61                    }
62                }
63                iter[v] += 1;
64            }
65            0
66        }
67
68        fn bfs(&self, s: usize) -> Vec<i32> {
69            let v = self.g.len();
70            let mut level = vec![-1; v];
71            level[s] = 0;
72            let mut deque = std::collections::VecDeque::new();
73            deque.push_back(s);
74            while let Some(v) = deque.pop_front() {
75                for e in self.g[v].iter() {
76                    if e.cap > 0 && level[e.to] < 0 {
77                        level[e.to] = level[v] + 1;
78                        deque.push_back(e.to);
79                    }
80                }
81            }
82            level
83        }
84
85        pub fn max_flow(&mut self, s: usize, t: usize) -> i64 {
86            let v = self.g.len();
87            let mut flow: i64 = 0;
88            loop {
89                let level = self.bfs(s);
90                if level[t] < 0 {
91                    return flow;
92                }
93                let mut iter = vec![0; v];
94                loop {
95                    let f = self.dfs(s, t, INF, &level, &mut iter);
96                    if f == 0 {
97                        break;
98                    }
99                    flow += f;
100                }
101            }
102        }
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::utils::test_helper::Tester;
110
111    /// Solve http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_6_A
112    #[test]
113    fn solve_grl_6_a() {
114        let tester = Tester::new("./assets/GRL_6_A/in/", "./assets/GRL_6_A/out/");
115        tester.test_solution(|sc| {
116            let v = sc.read();
117            let e = sc.read();
118
119            let mut dinitz = dinitz::Dinitz::new(v);
120            for _ in 0..e {
121                let from = sc.read();
122                let to = sc.read();
123                let c = sc.read();
124                dinitz.add_edge(from, to, c);
125            }
126            let flow = dinitz.max_flow(0, v - 1);
127            sc.write(format!("{}\n", flow));
128        });
129    }
130}