competitive_programming_rs/graph/
min_cost_flow.rs

1pub mod primal_dual {
2    use std::cmp;
3    use std::collections::BinaryHeap;
4    type Flow = i64;
5    type Cost = i64;
6    const INF: Cost = 1 << 60;
7    struct Edge {
8        to: usize,
9        capacity: Flow,
10        flow: Flow,
11        cost: Cost,
12        reverse_to: usize,
13    }
14    impl Edge {
15        fn residue(&self) -> Flow {
16            self.capacity - self.flow
17        }
18    }
19
20    pub struct MinimumCostFlowSolver {
21        graph: Vec<Vec<Edge>>,
22        previous_edge: Vec<(usize, usize)>,
23    }
24
25    impl MinimumCostFlowSolver {
26        pub fn new(n: usize) -> Self {
27            MinimumCostFlowSolver {
28                graph: (0..n).map(|_| Vec::new()).collect(),
29                previous_edge: vec![(0, 0); n],
30            }
31        }
32
33        pub fn add_edge(&mut self, from: usize, to: usize, capacity: Flow, cost: Cost) {
34            let reverse_from = self.graph[to].len();
35            let reverse_to = self.graph[from].len();
36            self.graph[from].push(Edge {
37                to,
38                capacity,
39                flow: 0,
40                cost,
41                reverse_to: reverse_from,
42            });
43            self.graph[to].push(Edge {
44                to: from,
45                capacity,
46                flow: capacity,
47                cost: -cost,
48                reverse_to,
49            });
50        }
51
52        /// Find the minimum cost to send `flow` through a flow network from `source` to `sink`.
53        pub fn solve(&mut self, source: usize, sink: usize, mut flow: Flow) -> Option<Flow> {
54            let n = self.graph.len();
55            let mut result = 0;
56            let mut h = vec![0; n];
57            let mut q: BinaryHeap<(Cost, usize)> = BinaryHeap::new();
58            while flow > 0 {
59                let mut dist = vec![INF; n];
60                dist[source] = 0;
61                q.push((0, source));
62                while let Some((current_dist, v)) = q.pop() {
63                    if dist[v] < current_dist {
64                        continue;
65                    }
66                    for (i, e) in self.graph[v].iter().enumerate() {
67                        if e.residue() == 0 {
68                            continue;
69                        }
70                        if dist[e.to] + h[e.to] > current_dist + h[v] + e.cost {
71                            dist[e.to] = current_dist + h[v] + e.cost - h[e.to];
72                            self.previous_edge[e.to] = (v, i);
73                            q.push((dist[e.to], e.to));
74                        }
75                    }
76                }
77
78                if dist[sink] == INF {
79                    return None;
80                }
81
82                for i in 0..n {
83                    h[i] += dist[i];
84                }
85                let mut df = flow;
86                let mut v = sink;
87                while v != source {
88                    let (prev_v, prev_e) = self.previous_edge[v];
89                    df = cmp::min(df, self.graph[prev_v][prev_e].residue());
90                    v = prev_v;
91                }
92                flow -= df;
93                result += df * h[sink];
94                let mut v = sink;
95                while v != source {
96                    let (prev_v, prev_e) = self.previous_edge[v];
97                    self.graph[prev_v][prev_e].flow += df;
98                    let reversed_edge_id = self.graph[prev_v][prev_e].reverse_to;
99                    self.graph[v][reversed_edge_id].flow -= df;
100                    v = prev_v;
101                }
102            }
103            Some(result)
104        }
105
106        /// Find the minimum cost to send `flow` through a flow network, which contains edges of
107        /// negative cost, from `source` to `sink`.
108        pub fn neg_solve(&mut self, source: usize, sink: usize, mut flow: Flow) -> Option<Flow> {
109            let n = self.graph.len();
110            let mut result = 0;
111            while flow > 0 {
112                let mut dist = vec![INF; n];
113                dist[source] = 0;
114                loop {
115                    let mut updated = false;
116                    for v in 0..n {
117                        if dist[v] == INF {
118                            continue;
119                        }
120
121                        for (i, e) in self.graph[v].iter().enumerate() {
122                            if e.residue() == 0 {
123                                continue;
124                            }
125                            if dist[e.to] > dist[v] + e.cost {
126                                dist[e.to] = dist[v] + e.cost;
127                                self.previous_edge[e.to] = (v, i);
128                                updated = true;
129                            }
130                        }
131                    }
132                    if !updated {
133                        break;
134                    }
135                }
136
137                if dist[sink] == INF {
138                    return None;
139                }
140
141                let mut df = flow;
142                let mut v = sink;
143                while v != source {
144                    let (prev_v, prev_e) = self.previous_edge[v];
145                    df = cmp::min(df, self.graph[prev_v][prev_e].residue());
146                    v = prev_v;
147                }
148                flow -= df;
149                result += df * dist[sink];
150                let mut v = sink;
151                while v != source {
152                    let (prev_v, prev_e) = self.previous_edge[v];
153                    self.graph[prev_v][prev_e].flow += df;
154                    let reversed_edge_id = self.graph[prev_v][prev_e].reverse_to;
155                    self.graph[v][reversed_edge_id].flow -= df;
156                    v = prev_v;
157                }
158            }
159            Some(result)
160        }
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use crate::utils::test_helper::Tester;
168
169    #[test]
170    fn solve_grl_6_b() {
171        let tester = Tester::new("./assets/GRL_6_B/in/", "./assets/GRL_6_B/out/");
172        tester.test_solution(|sc| {
173            let v: usize = sc.read();
174            let e: usize = sc.read();
175            let f: i64 = sc.read();
176
177            let mut solver = primal_dual::MinimumCostFlowSolver::new(v);
178            for _ in 0..e {
179                let u: usize = sc.read();
180                let v: usize = sc.read();
181                let c: i64 = sc.read();
182                let d: i64 = sc.read();
183                solver.add_edge(u, v, c, d);
184            }
185
186            let ans = match solver.solve(0, v - 1, f) {
187                Some(flow) => flow,
188                None => -1,
189            };
190            sc.write(format!("{}\n", ans));
191        });
192    }
193
194    #[test]
195    fn solve_grl_6_b_negative() {
196        let tester = Tester::new("./assets/GRL_6_B/in/", "./assets/GRL_6_B/out/");
197        tester.test_solution(|sc| {
198            let v: usize = sc.read();
199            let e: usize = sc.read();
200            let f: i64 = sc.read();
201
202            let mut solver = primal_dual::MinimumCostFlowSolver::new(v);
203            for _ in 0..e {
204                let u: usize = sc.read();
205                let v: usize = sc.read();
206                let c: i64 = sc.read();
207                let d: i64 = sc.read();
208                solver.add_edge(u, v, c, d);
209            }
210
211            let ans = match solver.neg_solve(0, v - 1, f) {
212                Some(flow) => flow,
213                None => -1,
214            };
215            sc.write(format!("{}\n", ans));
216        });
217    }
218}
competitive_programming_rs/graph/min_cost_flow.rs

competitive_programming_rs/graph/
min_cost_flow.rs
