competitive/graph/
minimum_cost_flow.rs

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

competitive/graph/
minimum_cost_flow.rs
