1use crate::{Cost, Error, Result, SolverStats, SolverStatus};
7use std::collections::VecDeque;
8use std::time::Instant;
9
10#[derive(Debug, Clone)]
12pub struct FlowNetwork {
13 pub num_nodes: usize,
15 adj: Vec<Vec<usize>>,
17 edges: Vec<FlowEdge>,
19}
20
21#[derive(Debug, Clone, Copy)]
23struct FlowEdge {
24 to: usize,
26 capacity: i64,
28 cost: i64,
30 flow: i64,
32 rev: usize,
34}
35
36impl FlowNetwork {
37 pub fn new(num_nodes: usize) -> Self {
39 Self {
40 num_nodes,
41 adj: vec![Vec::new(); num_nodes],
42 edges: Vec::new(),
43 }
44 }
45
46 pub fn add_edge(&mut self, from: usize, to: usize, capacity: i64, cost: i64) {
48 let forward_idx = self.edges.len();
49 let reverse_idx = self.edges.len() + 1;
50
51 self.edges.push(FlowEdge {
53 to,
54 capacity,
55 cost,
56 flow: 0,
57 rev: reverse_idx,
58 });
59 self.adj[from].push(forward_idx);
60
61 self.edges.push(FlowEdge {
63 to: from,
64 capacity: 0, cost: -cost, flow: 0,
67 rev: forward_idx,
68 });
69 self.adj[to].push(reverse_idx);
70 }
71
72 pub fn add_edge_with_capacity(&mut self, from: usize, to: usize, capacity: i64) {
74 self.add_edge(from, to, capacity, 0);
75 }
76
77 fn residual(&self, edge_idx: usize) -> i64 {
79 self.edges[edge_idx].capacity - self.edges[edge_idx].flow
80 }
81
82 fn push_flow(&mut self, edge_idx: usize, amount: i64) {
84 self.edges[edge_idx].flow += amount;
85 let rev = self.edges[edge_idx].rev;
86 self.edges[rev].flow -= amount;
87 }
88}
89
90#[derive(Debug, Clone)]
96pub struct MaxFlowResult {
97 pub max_flow: i64,
99 pub edge_flows: Vec<i64>,
101 pub status: SolverStatus,
103 pub stats: SolverStats,
105}
106
107pub fn max_flow(network: &FlowNetwork, source: usize, sink: usize) -> Result<MaxFlowResult> {
111 if source >= network.num_nodes || sink >= network.num_nodes {
112 return Err(Error::invalid_input("source or sink out of range"));
113 }
114 if source == sink {
115 return Err(Error::invalid_input("source and sink must be different"));
116 }
117
118 let start = Instant::now();
119 let n = network.num_nodes;
120
121 let mut net = network.clone();
123
124 let mut height = vec![0usize; n];
126 let mut excess = vec![0i64; n];
127
128 let mut current = vec![0usize; n];
130
131 let mut active: VecDeque<usize> = VecDeque::new();
133 let mut in_queue = vec![false; n];
134
135 height[source] = n;
137
138 let source_edges: Vec<usize> = net.adj[source].clone();
140 for edge_idx in source_edges {
141 let cap = net.residual(edge_idx);
142 if cap > 0 {
143 let to = net.edges[edge_idx].to;
144 net.push_flow(edge_idx, cap);
145 excess[to] += cap;
146 excess[source] -= cap;
147
148 if to != sink && to != source && !in_queue[to] {
149 active.push_back(to);
150 in_queue[to] = true;
151 }
152 }
153 }
154
155 let mut iterations = 0;
156
157 while let Some(u) = active.pop_front() {
159 in_queue[u] = false;
160
161 let activated = discharge(
163 &mut net,
164 &mut height,
165 &mut excess,
166 &mut current,
167 u,
168 source,
169 sink,
170 );
171 iterations += 1;
172
173 for v in activated {
175 if !in_queue[v] {
176 active.push_back(v);
177 in_queue[v] = true;
178 }
179 }
180
181 if excess[u] > 0 && u != source && u != sink && !in_queue[u] {
183 active.push_back(u);
184 in_queue[u] = true;
185 }
186 }
187
188 let edge_flows: Vec<i64> = (0..net.edges.len())
190 .step_by(2)
191 .map(|i| net.edges[i].flow)
192 .collect();
193
194 let elapsed = start.elapsed().as_secs_f64();
195
196 Ok(MaxFlowResult {
197 max_flow: excess[sink],
198 edge_flows,
199 status: SolverStatus::Optimal,
200 stats: SolverStats {
201 solve_time_seconds: elapsed,
202 iterations,
203 objective_value: Some(excess[sink] as f64),
204 ..Default::default()
205 },
206 })
207}
208
209fn discharge(
212 net: &mut FlowNetwork,
213 height: &mut [usize],
214 excess: &mut [i64],
215 current: &mut [usize],
216 u: usize,
217 source: usize,
218 sink: usize,
219) -> Vec<usize> {
220 let mut activated = Vec::new();
221
222 while excess[u] > 0 {
223 if current[u] >= net.adj[u].len() {
224 relabel(net, height, u);
226 current[u] = 0;
227 } else {
228 let edge_idx = net.adj[u][current[u]];
229 let v = net.edges[edge_idx].to;
230 let residual = net.residual(edge_idx);
231
232 if residual > 0 && height[u] == height[v] + 1 {
233 let push_amount = excess[u].min(residual);
235 net.push_flow(edge_idx, push_amount);
236 excess[u] -= push_amount;
237
238 let was_zero = excess[v] == 0;
240 excess[v] += push_amount;
241
242 if was_zero && v != source && v != sink {
244 activated.push(v);
245 }
246 } else {
247 current[u] += 1;
248 }
249 }
250 }
251
252 activated
253}
254
255fn relabel(net: &FlowNetwork, height: &mut [usize], u: usize) {
257 let mut min_height = usize::MAX;
258
259 for &edge_idx in &net.adj[u] {
260 if net.residual(edge_idx) > 0 {
261 let v = net.edges[edge_idx].to;
262 min_height = min_height.min(height[v]);
263 }
264 }
265
266 if min_height < usize::MAX {
267 height[u] = min_height + 1;
268 }
269}
270
271#[derive(Debug, Clone)]
277pub struct MinCostFlowResult {
278 pub flow: i64,
280 pub cost: Cost,
282 pub edge_flows: Vec<i64>,
284 pub status: SolverStatus,
286 pub stats: SolverStats,
288}
289
290#[derive(Debug, Clone)]
292pub struct MinCostFlowProblem {
293 pub network: FlowNetwork,
295 pub supplies: Vec<i64>,
297}
298
299impl MinCostFlowProblem {
300 pub fn new(network: FlowNetwork, supplies: Vec<i64>) -> Result<Self> {
302 if supplies.len() != network.num_nodes {
303 return Err(Error::dimension_mismatch(network.num_nodes, supplies.len()));
304 }
305 let total: i64 = supplies.iter().sum();
307 if total != 0 {
308 return Err(Error::invalid_input(format!(
309 "supplies must sum to 0, got {}",
310 total
311 )));
312 }
313 Ok(Self { network, supplies })
314 }
315
316 pub fn source_sink(
318 network: FlowNetwork,
319 source: usize,
320 sink: usize,
321 flow_demand: i64,
322 ) -> Result<Self> {
323 let mut supplies = vec![0i64; network.num_nodes];
324 supplies[source] = flow_demand;
325 supplies[sink] = -flow_demand;
326 Self::new(network, supplies)
327 }
328}
329
330pub fn min_cost_flow(problem: &MinCostFlowProblem) -> Result<MinCostFlowResult> {
335 let start = Instant::now();
336 let n = problem.network.num_nodes;
337
338 let mut net = problem.network.clone();
340 let mut supply = problem.supplies.clone();
341
342 let mut total_flow: i64 = 0;
343 let mut total_cost: Cost = 0;
344 let mut iterations = 0;
345
346 loop {
348 iterations += 1;
349
350 let source = supply.iter().position(|&s| s > 0);
352 let sink = supply.iter().position(|&s| s < 0);
353
354 match (source, sink) {
355 (Some(s), Some(t)) => {
356 match bellman_ford_path(&net, s, t) {
358 Some((path, path_cost)) => {
359 let mut bottleneck = supply[s].min(-supply[t]);
361 for &edge_idx in &path {
362 bottleneck = bottleneck.min(net.residual(edge_idx));
363 }
364
365 if bottleneck <= 0 {
366 break; }
368
369 for &edge_idx in &path {
371 net.push_flow(edge_idx, bottleneck);
372 }
373
374 supply[s] -= bottleneck;
375 supply[t] += bottleneck;
376 total_flow += bottleneck;
377 total_cost += bottleneck * path_cost;
378 }
379 None => {
380 if supply.iter().any(|&s| s != 0) {
382 return Err(Error::infeasible(
383 "no augmenting path but unsatisfied supply/demand",
384 ));
385 }
386 break;
387 }
388 }
389 }
390 _ => break, }
392
393 if iterations > n * net.edges.len() * 1000 {
395 return Err(Error::NoConvergence { iterations });
396 }
397 }
398
399 if supply.iter().any(|&s| s != 0) {
401 return Err(Error::infeasible("could not satisfy all supplies/demands"));
402 }
403
404 let edge_flows: Vec<i64> = (0..net.edges.len())
406 .step_by(2)
407 .map(|i| net.edges[i].flow)
408 .collect();
409
410 let elapsed = start.elapsed().as_secs_f64();
411
412 Ok(MinCostFlowResult {
413 flow: total_flow,
414 cost: total_cost,
415 edge_flows,
416 status: SolverStatus::Optimal,
417 stats: SolverStats {
418 solve_time_seconds: elapsed,
419 iterations,
420 objective_value: Some(total_cost as f64),
421 ..Default::default()
422 },
423 })
424}
425
426fn bellman_ford_path(net: &FlowNetwork, source: usize, sink: usize) -> Option<(Vec<usize>, Cost)> {
429 let n = net.num_nodes;
430 let mut dist = vec![i64::MAX; n];
431 let mut parent_edge: Vec<Option<usize>> = vec![None; n];
432
433 dist[source] = 0;
434
435 for _ in 0..n {
437 let mut changed = false;
438 for u in 0..n {
439 if dist[u] == i64::MAX {
440 continue;
441 }
442 for &edge_idx in &net.adj[u] {
443 let edge = &net.edges[edge_idx];
444 if net.residual(edge_idx) > 0 {
445 let new_dist = dist[u].saturating_add(edge.cost);
446 if new_dist < dist[edge.to] {
447 dist[edge.to] = new_dist;
448 parent_edge[edge.to] = Some(edge_idx);
449 changed = true;
450 }
451 }
452 }
453 }
454 if !changed {
455 break;
456 }
457 }
458
459 if dist[sink] == i64::MAX {
460 return None;
461 }
462
463 let mut path = Vec::new();
465 let mut current = sink;
466 while let Some(edge_idx) = parent_edge[current] {
467 path.push(edge_idx);
468 let rev_idx = net.edges[edge_idx].rev;
470 current = net.edges[rev_idx].to;
471 if current == source {
472 break;
473 }
474 }
475 path.reverse();
476
477 Some((path, dist[sink]))
478}
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483
484 #[test]
485 fn test_simple_max_flow() {
486 let mut net = FlowNetwork::new(5);
512 net.add_edge_with_capacity(0, 1, 10);
513 net.add_edge_with_capacity(0, 2, 10);
514 net.add_edge_with_capacity(1, 2, 2);
515 net.add_edge_with_capacity(1, 3, 4);
516 net.add_edge_with_capacity(1, 4, 8);
517 net.add_edge_with_capacity(2, 4, 9);
518 net.add_edge_with_capacity(3, 4, 10);
519
520 let result = max_flow(&net, 0, 4).unwrap();
521 assert!(result.max_flow >= 17, "max_flow was {}", result.max_flow);
538 assert!(result.max_flow <= 20, "max_flow was {}", result.max_flow);
540 }
541
542 #[test]
543 fn test_max_flow_simple_path() {
544 let mut net = FlowNetwork::new(3);
546 net.add_edge_with_capacity(0, 1, 5);
547 net.add_edge_with_capacity(1, 2, 3);
548
549 let result = max_flow(&net, 0, 2).unwrap();
550 assert_eq!(result.max_flow, 3); }
552
553 #[test]
554 fn test_max_flow_parallel_paths() {
555 let mut net = FlowNetwork::new(4);
557 net.add_edge_with_capacity(0, 1, 10);
558 net.add_edge_with_capacity(1, 3, 10);
559 net.add_edge_with_capacity(0, 2, 10);
560 net.add_edge_with_capacity(2, 3, 10);
561
562 let result = max_flow(&net, 0, 3).unwrap();
563 assert_eq!(result.max_flow, 20); }
565
566 #[test]
567 fn test_min_cost_flow_simple() {
568 let mut net = FlowNetwork::new(4);
577 net.add_edge(0, 1, 10, 1);
578 net.add_edge(0, 2, 10, 5);
579 net.add_edge(1, 3, 10, 1);
580 net.add_edge(2, 3, 10, 1);
581
582 let problem = MinCostFlowProblem::source_sink(net, 0, 3, 5).unwrap();
583 let result = min_cost_flow(&problem).unwrap();
584
585 assert_eq!(result.flow, 5);
586 assert_eq!(result.cost, 10); }
588
589 #[test]
590 fn test_min_cost_flow_with_supplies() {
591 let mut net = FlowNetwork::new(4);
599 net.add_edge(0, 2, 10, 1);
600 net.add_edge(0, 3, 10, 3);
601 net.add_edge(1, 2, 10, 2);
602 net.add_edge(1, 3, 10, 1);
603
604 let supplies = vec![5, 5, -5, -5];
605 let problem = MinCostFlowProblem::new(net, supplies).unwrap();
606 let result = min_cost_flow(&problem).unwrap();
607
608 assert_eq!(result.flow, 10); assert_eq!(result.cost, 10);
611 }
612
613 #[test]
614 fn test_infeasible_supplies() {
615 let net = FlowNetwork::new(2);
616 let supplies = vec![5, 0]; let result = MinCostFlowProblem::new(net, supplies);
619 assert!(result.is_err());
620 }
621}