competitive/graph/
minimum_cost_flow.rs1use super::BidirectionalSparseGraph;
2
3#[derive(Debug, Clone)]
4pub struct PrimalDualBuilder {
5 vsize: usize,
6 edges: Vec<(usize, usize)>,
7 capacities: Vec<u64>,
8 costs: Vec<i64>,
9 has_negedge: bool,
10}
11impl PrimalDualBuilder {
12 pub fn new(vsize: usize, esize_expect: usize) -> Self {
13 Self {
14 vsize,
15 edges: Vec::with_capacity(esize_expect),
16 capacities: Vec::with_capacity(esize_expect * 2),
17 costs: Vec::with_capacity(esize_expect * 2),
18 has_negedge: false,
19 }
20 }
21 pub fn add_edge(&mut self, from: usize, to: usize, cap: u64, cost: i64) {
22 self.edges.push((from, to));
23 self.capacities.push(cap);
24 self.capacities.push(0);
25 self.has_negedge |= cost < 0;
26 self.costs.push(cost);
27 self.costs.push(-cost);
28 }
29 pub fn gen_graph(&mut self) -> BidirectionalSparseGraph {
30 let edges = std::mem::take(&mut self.edges);
31 BidirectionalSparseGraph::from_edges(self.vsize, edges)
32 }
33 pub fn build(self, graph: &BidirectionalSparseGraph) -> PrimalDual<'_> {
34 let PrimalDualBuilder {
35 vsize,
36 capacities,
37 costs,
38 has_negedge,
39 ..
40 } = self;
41 PrimalDual {
42 graph,
43 capacities,
44 costs,
45 potential: std::iter::repeat(0).take(vsize).collect(),
46 dist: Vec::with_capacity(vsize),
47 prev_vertex: std::iter::repeat(0).take(vsize).collect(),
48 prev_edge: std::iter::repeat(0).take(vsize).collect(),
49 has_negedge,
50 }
51 }
52}
53impl Extend<(usize, usize, u64, i64)> for PrimalDualBuilder {
54 fn extend<T: IntoIterator<Item = (usize, usize, u64, i64)>>(&mut self, iter: T) {
55 for (from, to, cap, cost) in iter {
56 self.add_edge(from, to, cap, cost)
57 }
58 }
59}
60
61#[derive(Debug)]
62pub struct PrimalDual<'a> {
63 graph: &'a BidirectionalSparseGraph,
64 capacities: Vec<u64>,
65 costs: Vec<i64>,
66 potential: Vec<i64>,
67 dist: Vec<i64>,
68 prev_vertex: Vec<usize>,
69 prev_edge: Vec<usize>,
70 has_negedge: bool,
71}
72impl PrimalDual<'_> {
73 pub fn builder(vsize: usize, esize_expect: usize) -> PrimalDualBuilder {
74 PrimalDualBuilder::new(vsize, esize_expect)
75 }
76 fn bellman_ford(&mut self, s: usize) {
77 self.potential.clear();
78 self.potential.resize(self.graph.vertices_size(), i64::MAX);
79 self.potential[s] = 0;
80 for _ in 1..self.graph.vertices_size() {
81 let mut end = true;
82 for u in self.graph.vertices() {
83 for a in self.graph.adjacencies(u) {
84 let ncost = self.potential[u].saturating_add(self.costs[a.id]);
85 if self.capacities[a.id] > 0 && self.potential[a.to] > ncost {
86 self.potential[a.to] = ncost;
87 end = false;
88 }
89 }
90 }
91 if end {
92 break;
93 }
94 }
95 }
96 fn dijkstra(&mut self, s: usize, t: usize) -> bool {
97 use std::{cmp::Reverse, collections::BinaryHeap};
98 self.dist.clear();
99 self.dist.resize(self.graph.vertices_size(), i64::MAX);
100 self.dist[s] = 0;
101 let mut heap = BinaryHeap::new();
102 heap.push((Reverse(0), s));
103 while let Some((Reverse(d), u)) = heap.pop() {
104 if self.dist[u] < d {
105 continue;
106 }
107 if !self.has_negedge && u == t {
108 break;
109 }
110 for a in self.graph.adjacencies(u) {
111 let ncost = (self.dist[u].saturating_add(self.costs[a.id]))
112 .saturating_add(self.potential[u].saturating_sub(self.potential[a.to]));
113 if self.capacities[a.id] > 0 && self.dist[a.to] > ncost {
114 self.dist[a.to] = ncost;
115 self.prev_vertex[a.to] = u;
116 self.prev_edge[a.to] = a.id;
117 heap.push((Reverse(ncost), a.to));
118 }
119 }
120 }
121 self.dist[t] != i64::MAX
122 }
123 pub fn minimum_cost_flow_limited(&mut self, s: usize, t: usize, limit: u64) -> (u64, i64) {
125 let mut flow = 0;
126 let mut cost = 0;
127 if self.has_negedge {
128 self.bellman_ford(s);
129 }
130 while flow < limit && self.dijkstra(s, t) {
131 for (p, d) in self.potential.iter_mut().zip(self.dist.iter()) {
132 *p = p.saturating_add(*d);
133 }
134 let mut f = limit - flow;
135 let mut v = t;
136 while v != s {
137 f = f.min(self.capacities[self.prev_edge[v]]);
138 v = self.prev_vertex[v];
139 }
140 flow += f;
141 cost += f as i64 * self.potential[t];
142 let mut v = t;
143 while v != s {
144 self.capacities[self.prev_edge[v]] -= f;
145 self.capacities[self.prev_edge[v] ^ 1] += f;
146 v = self.prev_vertex[v];
147 }
148 }
149 (flow, cost)
150 }
151 pub fn minimum_cost_flow(&mut self, s: usize, t: usize) -> (u64, i64) {
153 self.minimum_cost_flow_limited(s, t, u64::MAX)
154 }
155 pub fn get_flow(&self, eid: usize) -> u64 {
156 self.capacities[eid * 2 + 1]
157 }
158}